Upload test_model.py with huggingface_hub

test_model.py

@@ -1,24 +1,30 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 import torch
 import json
 
-# Detect device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Device set to use: {device}")
 
 # Load base model and tokenizer
-base_model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-adapter_repo = "Harish2002/cli-lora-tinyllama"
+base_model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0").to(device)
+tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 
-tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-base_model = AutoModelForCausalLM.from_pretrained(base_model_name)
-model = PeftModel.from_pretrained(base_model, adapter_repo)
-model = model.to(device)
+# Load LoRA adapter
+model = PeftModel.from_pretrained(base_model, "Harish2002/cli-lora-tinyllama")
+model.to(device)
 model.eval()
 
-# Test prompts
-test_prompts = {
+# Utility function to generate answers
+def generate_answer(question):
+    prompt = f"{question}\nAnswer:"
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=128)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True).replace(prompt, "").strip()
+
+# Questions to test
+questions = {
     "Git": "How do I create a new branch and switch to it in Git?",
     "Bash": "How to list all files including hidden ones?",
     "Grep": "How do I search for a pattern in multiple files using grep?",
@@ -26,21 +32,18 @@ test_prompts = {
     "Python venv": "How do I activate a virtual environment on Windows?"
 }
 
-# Run test and store results
+# Run test and save results
 results = {}
-for topic, prompt in test_prompts.items():
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    with torch.no_grad():
-        outputs = model.generate(**inputs, max_new_tokens=128)
-    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    results[topic] = {
-        "question": prompt,
-        "answer": answer
-    }
-    print(f"\n🧪 {topic}:\nQ: {prompt}\nA: {answer}")
-
-# Save to file
-with open("test_outputs.json", "w", encoding="utf-8") as f:
-    json.dump(results, f, indent=4)
+
+for category, question in questions.items():
+    print(f"\n🧪 {category}:")
+    print(f"Q: {question}")
+    answer = generate_answer(question)
+    print(f"A: {answer}\n")
+    results[category] = {"question": question, "answer": answer}
+
+# Save to JSON
+with open("test_outputs.json", "w") as f:
+    json.dump(results, f, indent=2)
 
 print("\n✅ All outputs saved to test_outputs.json")