simple_eval.py

"""
Evaluation script for Llama 3.2 + LoRA adapters on BoolQ-style dataset.
FIXED VERSION: Handles prompt mismatch, better prediction parsing, and metrics
"""

import os
import json
import argparse
import torch
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel

# ============================================================================ #
# 1. Parse arguments
# ============================================================================ #
parser = argparse.ArgumentParser()
parser.add_argument("--adapter_path", type=str, required=True,
                    help="Path to adapter (e.g. ./llama3-boolq-qlora/adapter_correct)")
parser.add_argument("--data_path", type=str, required=True,
                    help="Path to evaluation dataset (JSONL format)")
parser.add_argument("--output_path", type=str, default="eval_results.json",
                    help="File to save predictions and metrics")
parser.add_argument("--base_model", type=str, default="meta-llama/Llama-3.2-1B",
                    help="Base model name")
parser.add_argument("--max_new_tokens", type=int, default=5,
                    help="Max tokens to generate (we only need Yes/No)")
parser.add_argument("--load_in_4bit", action="store_true",
                    help="Load model in 4-bit for memory efficiency")
parser.add_argument("--batch_size", type=int, default=1,
                    help="Batch size for evaluation (set to 1 for now)")
args = parser.parse_args()

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {DEVICE}")
# ============================================================================ #
# 2. Load model + adapter
# ============================================================================ #
print(f"Loading base model: {args.base_model}")

# CHANGE 1: Support 4-bit loading for memory efficiency
if args.load_in_4bit:
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )
    base_model = AutoModelForCausalLM.from_pretrained(
        args.base_model,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True
    )
else:
    base_model = AutoModelForCausalLM.from_pretrained(
        args.base_model,
        device_map="auto",
        torch_dtype=torch.bfloat16
    )

tokenizer = AutoTokenizer.from_pretrained(args.base_model)
# CHANGE 2: Ensure padding token is set (critical for batch processing)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"  # Left padding for generation

print(f"Loading adapter: {args.adapter_path}")
model = PeftModel.from_pretrained(base_model, args.adapter_path)
# model = base_model
model.eval()

# ============================================================================ #
# 3. Load dataset
# ============================================================================ #
def load_jsonl(path):
    with open(path, "r", encoding="utf-8") as f:
        return [json.loads(line) for line in f if line.strip()]

data = load_jsonl(args.data_path)
print(f"Loaded {len(data)} samples for evaluation")

# ============================================================================ #
# 4. Format prompt (MUST MATCH TRAINING)
# ============================================================================ #
def format_eval_prompt(example):
    """
    CHANGE 3: This MUST match the training prompt format exactly
    Your training used: "Passage: ... Question: ... Answer (Yes or No):"
    """
    prompt = (
        f"Passage: {example['passage']}\n\n"
        f"Question: {example['question']}\n\n"
        f"Answer (Yes or No):"
    )
    return prompt

# ============================================================================ #
# 5. Prediction parsing
# ============================================================================ #
def extract_prediction(generated_text):
    """
    CHANGE 4: More robust prediction extraction
    Handles: "Yes", "yes", " Yes", "Yes.", "Yes,", etc.
    """
    text = generated_text.strip().lower()
    
    # Direct match
    if text.startswith('yes'):
        return "Yes"
    elif text.startswith('no'):
        return "No"
    
    # Check first word
    words = text.split()
    if words:
        first_word = words[0].rstrip('.,!?;:')
        if first_word == 'yes':
            return "Yes"
        elif first_word == 'no':
            return "No"
    
    # Fallback: search for yes/no anywhere in response
    if 'yes' in text and 'no' not in text:
        return "Yes"
    elif 'no' in text and 'yes' not in text:
        return "No"
    
    # CHANGE 5: Return None for unclear predictions (don't default to "No")
    return None

# ============================================================================ #
# 6. Evaluation loop
# ============================================================================ #
correct = 0
unclear = 0  # Track unclear predictions
results = []

print("\nStarting evaluation...")
for ex in tqdm(data, desc="Evaluating"):
    gold_answer = "Yes" if ex["answer"] else "No"
    
    # Format prompt (must match training)
    prompt = format_eval_prompt(ex)
    
    # CHANGE 6: Add return tensors and move to device properly
    inputs = tokenizer(prompt, return_tensors="pt", padding=True)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    
    # Generate
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=args.max_new_tokens,
            do_sample=False,  # Greedy decoding for consistency
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    
    # CHANGE 7: Decode only the generated part (not the prompt)
    prompt_length = inputs['input_ids'].shape[1]
    generated_ids = outputs[0][prompt_length:]
    generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
    
    # Extract prediction
    pred = extract_prediction(generated_text)
    
    # CHANGE 8: Track unclear predictions separately
    if pred is None:
        unclear += 1
        pred = "Unclear"
        is_correct = False
    else:
        is_correct = (pred == gold_answer)
        if is_correct:
            correct += 1
    
    # CHANGE 9: Store more detailed results
    results.append({
        "question": ex["question"],
        "passage": ex["passage"][:200] + "...",  # Truncate for readability
        "gold": gold_answer,
        "pred": pred,
        "generated_text": generated_text,
        "correct": is_correct
    })

# ============================================================================ #
# 7. Calculate and display metrics
# ============================================================================ #
total = len(data)
accuracy = correct / total
clarity_rate = (total - unclear) / total

print(f"\n{'='*60}")
print(f"EVALUATION RESULTS")
print(f"{'='*60}")
print(f"Adapter: {args.adapter_path}")
print(f"Dataset: {args.data_path}")
print(f"{'='*60}")
print(f"Total samples: {total}")
print(f"Correct: {correct}")
print(f"Incorrect: {total - correct - unclear}")
print(f"Unclear: {unclear}")
print(f"{'='*60}")
print(f"Accuracy: {accuracy*100:.2f}% ({correct}/{total})")
print(f"Clarity rate: {clarity_rate*100:.2f}%")
print(f"{'='*60}")

# CHANGE 10: Show some example predictions
print("\nSample predictions (first 3):")
for i, r in enumerate(results[:3]):
    print(f"\n[{i+1}] Question: {r['question']}")
    print(f"    Gold: {r['gold']} | Pred: {r['pred']} | Generated: '{r['generated_text']}'")

# ============================================================================ #
# 8. Save results with enhanced metrics
# ============================================================================ #
output = {
    "metadata": {
        "adapter_path": args.adapter_path,
        "data_path": args.data_path,
        "base_model": args.base_model,
        "max_new_tokens": args.max_new_tokens,
    },
    "metrics": {
        "accuracy": accuracy,
        "num_correct": correct,
        "num_samples": total,
        "num_unclear": unclear,
        "clarity_rate": clarity_rate,
    },
    "predictions": results
}

os.makedirs(os.path.dirname(args.output_path) or ".", exist_ok=True)
with open(args.output_path, "w", encoding="utf-8") as f:
    json.dump(output, f, indent=2, ensure_ascii=False)

print(f"\n✅ Results saved to {args.output_path}")

# CHANGE 11: Also save a summary CSV for easy comparison
import csv
summary_path = args.output_path.replace('.json', '_summary.csv')
with open(summary_path, 'w', newline='') as f:
    writer = csv.writer(f)
    writer.writerow(['adapter_path', 'accuracy', 'correct', 'total', 'unclear'])
    writer.writerow([args.adapter_path, f"{accuracy:.4f}", correct, total, unclear])

print(f"Summary saved to {summary_path}")