#!/usr/bin/env python3
# ============================================================
# ZabaanAI-v2: Evaluation Script
# Tasks: MT-Bench, PMLU, grammar error detection, translation
# ============================================================
import os, warnings, json
warnings.filterwarnings('ignore')

from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
from datasets import load_dataset
import evaluate

MODEL_PATH = '/app/output/zabaanai-v2-sft/final'
HF_REPO    = 'shaikhsalman/zabaanai-v2-sft'

print('Loading model...')
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
base = AutoModelForCausalLM.from_pretrained(
    'Qwen/Qwen2.5-7B-Instruct', torch_dtype='bfloat16', device_map='auto', trust_remote_code=True
)
model = PeftModel.from_pretrained(base, MODEL_PATH)

generator = pipeline(
    'text-generation', model=model, tokenizer=tokenizer,
    max_new_tokens=512, temperature=0.7, top_p=0.9
)

# ── Test Prompts (Pakistan-focused) ─────────────────────────
test_prompts = [
    # Urdu
    {'lang': 'ur', 'prompt': '<|im_start|>user\nپاکستان کے صوبے کون کون سے ہیں؟<|im_end|>\n<|im_start|>assistant\n'},
    {'lang': 'ur', 'prompt': '<|im_start|>user\nاردو میں ایک مختصر نظم لکھیں۔<|im_end|>\n<|im_start|>assistant\n'},
    # Punjabi (Shahmukhi)
    {'lang': 'pa', 'prompt': '<|im_start|>user\nپنجابی وچ تاریخ دیاں اہم تریخاں لکھو۔<|im_end|>\n<|im_start|>assistant\n'},
    # Sindhi
    {'lang': 'sd', 'prompt': '<|im_start|>user\nسنڌي ۾ ڳالهه ٻولهه لکھو ته پاڪستان جي آزادي بابت۔<|im_end|>\n<|im_start|>assistant\n'},
    # Pashto
    {'lang': 'ps', 'prompt': '<|im_start|>user\nپښتو کې د پاکستان په اړه لنډه مقاله ولیکه۔<|im_end|>\n<|im_start|>assistant\n'},
    # English
    {'lang': 'en', 'prompt': '<|im_start|>user\nExplain the education system in Pakistan up to FSC level.<|im_end|>\n<|im_start|>assistant\n'},
    # Roman Urdu
    {'lang': 'rom-ur', 'prompt': '<|im_start|>user\nRoman Urdu mein Pakistan ki q抵mat ke baare mein batao.<|im_end|>\n<|im_start|>assistant\n'},
    # Code-mixed
    {'lang': 'mixed', 'prompt': '<|im_start|>user\nExplain CSS exam preparation strategy in Urdu.<|im_end|>\n<|im_start|>assistant\n'},
]

results = []
print('Running evaluation...')
for item in test_prompts:
    out = generator(item['prompt'], return_full_text=False)
    generated = out[0]['generated_text'] if out else ''
    results.append({
        'language': item['lang'],
        'prompt': item['prompt'].replace('<|im_start|>','').replace('<|im_end|>',''),
        'generated': generated.strip(),
        'length': len(generated.split()),
    })
    print(f'  [{item[\"lang\"]}] → {len(generated.split())} words')

# Save results
out_path = '/app/data/evaluation_results.json'
os.makedirs('/app/data', exist_ok=True)
with open(out_path, 'w', encoding='utf-8') as f:
    json.dump(results, f, indent=2, ensure_ascii=False)

print(f'\nResults saved to {out_path}')
print('\nSample outputs:')
for r in results[:3]:
    print(f'\n[{r[\"language\"]}] Prompt: {r[\"prompt\"][:80]}')
    print(f'→ Response: {r[\"generated\"][:200]}')