import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset
from nltk.translate.bleu_score import corpus_bleu

# ----------------------------
# Configuration
# ----------------------------
MODEL_DIR = "./qwen2.5-7b-uk-fr-lora"  # dossier où tu as sauvegardé LoRA
VALIDATION_FILE = "validation.jsonl"  # petit subset de test (5-50 phrases)
MAX_INPUT_LENGTH = 1024
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

print("=== Loading model and tokenizer ===")
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.model_max_length = MAX_INPUT_LENGTH

model = AutoModelForCausalLM.from_pretrained(
    MODEL_DIR,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True
)
model.eval()

print("Model loaded.")

# ----------------------------
# Load validation dataset
# ----------------------------
print("Loading validation dataset...")
dataset = load_dataset("json", data_files=VALIDATION_FILE)
examples = dataset["train"]  # petit subset
print(f"{len(examples)} examples loaded for testing.")

# ----------------------------
# Function to generate translation
# ----------------------------
def translate(text):
    prompt = f"Translate the following Ukrainian text into French:\nUkrainian: {text}\nFrench:"
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=MAX_INPUT_LENGTH).to(DEVICE)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            do_sample=False,  # deterministic
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id
        )
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Remove prompt from result
    return result.replace(prompt, "").strip()

# ----------------------------
# Test all examples and compute BLEU
# ----------------------------
print("Generating translations...")
references = []
hypotheses = []

for i, example in enumerate(examples):
    src_text = example["text"]
    ref_text = example["translation"]
    pred_text = translate(src_text)
    
    print(f"\n[{i+1}] Source: {src_text}")
    print(f"    Reference: {ref_text}")
    print(f"    Prediction: {pred_text}")

    # Prepare for BLEU (tokenized by space)
    references.append([ref_text.split()])
    hypotheses.append(pred_text.split())

# Compute corpus BLEU
bleu_score = corpus_bleu(references, hypotheses)
print(f"\n=== Corpus BLEU score: {bleu_score:.4f} ===")