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