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

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device set to use: {device}")

# Load base model and tokenizer
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")

# Load LoRA adapter
model = PeftModel.from_pretrained(base_model, "Harish2002/cli-lora-tinyllama")
model.to(device)
model.eval()

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

# Run test and save results
results = {}

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")