import sys
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from pathlib import Path

# Fix Windows encoding issues for Hebrew characters
if sys.platform == "win32":
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

MODEL_DIR = Path("output/hebrew-coder-1.7b")

def test_model():
    print(f"Loading model from {MODEL_DIR.resolve()}...")
    
    try:
        tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_DIR,
            device_map="auto",
            torch_dtype="auto",
            trust_remote_code=True
        )
        print("Model loaded successfully!")
        
        prompts = [
            "כתוב פונקציה בפייתון שמחשבת עצרת של מספר",
            "Write a Python function to check if a number is prime, and explain it in Hebrew."
        ]
        
        for prompt in prompts:
            print(f"\nPrompt: {prompt}")
            
            messages = [
                {"role": "user", "content": prompt}
            ]
            
            # Use the proper instruct template so the model knows to respond
            inputs = tokenizer.apply_chat_template(
                messages, 
                add_generation_prompt=True, 
                return_tensors="pt"
            ).to(model.device)
            
            with torch.no_grad():
                outputs = model.generate(
                    inputs,
                    max_new_tokens=200,
                    do_sample=True,
                    temperature=0.7,
                    top_p=0.9
                )
            
            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
            print("-" * 40)
            print(result)
            print("-" * 40)
            
    except Exception as e:
        print(f"Error testing model: {e}")

if __name__ == "__main__":
    test_model()
