# test_model.py - FINAL CORRECTED VERSION
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import os 

# --- 1. Define the path where the Trainer saved the model ---
# Based on your directory structure, checkpoint-500 is the most up-to-date save.
MODEL_PATH = "./finetuned_model/checkpoint-1250" 


# --- 2. Load the fine-tuned model and tokenizer ---
try:
    # Ensure the path exists before attempting to load
    if not os.path.exists(MODEL_PATH):
        raise FileNotFoundError(f"Directory not found: {MODEL_PATH}")

    # We load the tokenizer and the model from your local save directory
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    model = AutoModelForCausalLM.from_pretrained(MODEL_PATH)
    
    # Reapply the padding token fix to ensure generation works smoothly
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        
    print(f"✅ Model and Tokenizer loaded successfully from {MODEL_PATH}!")
    
except Exception as e:
    print(f"❌ Error loading model: {e}")
    print("Ensure the checkpoint folder exists and contains `config.json` and weights.")
    exit()

# --- 3. Define the prompt based on your formatting function ---
# The prompt must mimic the structure used during training.
instruction = ""
input_text = "Hello?"

# Recreate the exact prompt structure used during training, but stop before the "Output:"
prompt = f"{instruction}\nInput: {input_text}\nOutput: "


# --- 4. Tokenize the prompt ---
# We use the same tokenizer and ensure the output is a PyTorch tensor
inputs = tokenizer(prompt, return_tensors="pt")


# --- 5. Generate the output ---
# The generate method is key for text generation tasks (Causal LMs)
output = model.generate(
    **inputs,
    max_new_tokens=50,       # The maximum number of tokens to generate after the prompt
    num_beams=1,             # Use 1 for greedy decoding (faster/simpler)
    do_sample=True,          # Use sampling for more creative/diverse answers
    temperature=0.7,         # Controls randomness (lower is more deterministic)
    top_k=50,                # Limits sampling to the top 50 most likely tokens
    eos_token_id=tokenizer.eos_token_id, # Stop generation when the EOS token is predicted
)


# --- 6. Decode and Display the Result ---
# Decode the generated token IDs back into a readable string
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)

print("\n--- Model Output ---")
print(f"Prompt Sent:\n{prompt}")
print("\nGenerated Text (Full Sequence):")
print(generated_text)

# To see only the newly generated part (the answer), you can strip the prompt:
only_answer = generated_text.replace(prompt, "").strip()
print(f"\nExtracted Answer:\n{only_answer}")