test_model.py

#!/usr/bin/env python3
# Test script for the DeepSeek D&D LoRA model
# This script uses the Hugging Face Hub to track usage metrics

from huggingface_hub import snapshot_download
import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

print("Loading model and tokenizer...")

# Load the base model and tokenizer
base_model_id = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
model = AutoModelForCausalLM.from_pretrained(
    base_model_id, 
    torch_dtype=torch.float16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(base_model_id)

# Load the LoRA adapter - this will be tracked by HF Hub
adapter_model_id = "chendren/deepseek-dnd-lora"
model = PeftModel.from_pretrained(model, adapter_model_id)

# Test prompts
test_prompts = [
    "Create a D&D character with the following details: Race: Half-Elf, Class: Bard, Background: Entertainer",
    "Design a D&D adventure hook set in a dark forest",
    "Create a magical item for D&D 5e that would be suitable for a level 5 rogue",
    "Write a description for a fantasy tavern in a D&D setting"
]

for i, prompt in enumerate(test_prompts):
    print(f"\n==== Test Prompt {i+1} ====")
    print(prompt)
    print("\n==== Response ====")
    
    inputs = tokenizer(f"You are a Dungeons & Dragons assistant. {prompt}", return_tensors="pt").to(model.device)
    
    outputs = model.generate(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_new_tokens=500,
        temperature=0.7,
        top_p=0.9,
        top_k=50,
        repetition_penalty=1.1,
        do_sample=True
    )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(response)
    print("\n" + "="*50)