from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch
import json

# Detect device
device = torch.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"

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)
model.eval()

# Test prompts
test_prompts = {
    "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?",
    "Tar/Gzip": "How to extract a .tar.gz file?",
    "Python venv": "How do I activate a virtual environment on Windows?"
}

# Run test and store 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)

print("\n✅ All outputs saved to test_outputs.json")