import os import sys import json import argparse from typing import List, Dict, Tuple, Union, Any from pathlib import Path import pandas as pd from tqdm.auto import tqdm import torch from transformers import ( MBart50Tokenizer, #type: ignore MBartForConditionalGeneration, #type: ignore MT5ForConditionalGeneration, #type: ignore MT5TokenizerFast, #type: ignore ) from peft import PeftModel, PeftConfig import evaluate # Add parent directory to sys.path sys.path.append(str(Path(__file__).resolve().parent.parent)) from models.rule_based_mt import TransferBasedMT from models.statistical_mt import SMTExtended, LanguageModel # Device configuration DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Load configuration with open("config.json", "r") as json_file: CONFIG = json.load(json_file) def parse_arguments() -> argparse.Namespace: """Parse command-line arguments.""" parser = argparse.ArgumentParser(description="Evaluate English-Vietnamese Machine Translation Models") parser.add_argument("--test_file", type=str, default='data/test_cleaned_dataset.csv', help="Path to test CSV file") parser.add_argument("--output_dir", type=str, default="results", help="Directory to save results") return parser.parse_args() class ModelLoader: """Handles loading of translation models.""" @staticmethod def load_smt() -> None: """Load Statistical Machine Translation model.""" try: smt = SMTExtended() model_dir = "checkpoints" if os.path.exists(model_dir) and os.path.isfile(os.path.join(model_dir, "phrase_table.pkl")): print("Loading existing model...") smt.load_model() else: print("Training new smt...") stats = smt.train() print(f"Training complete: {stats}") print("SMT model loaded successfully!") return smt except Exception as e: raise RuntimeError(f"Failed to load SMT model: {str(e)}") @staticmethod def load_mbart50() -> Tuple[MBartForConditionalGeneration, MBart50Tokenizer]: """Load fine-tuned MBart50 model and tokenizer.""" try: model_config = CONFIG["mbart50"]["paths"] model = MBartForConditionalGeneration.from_pretrained(model_config["base_model_name"]) peft_config = PeftConfig.from_pretrained(model_config["checkpoint_path"]) model = PeftModel.from_pretrained(model, model_config["checkpoint_path"]) tokenizer = MBart50Tokenizer.from_pretrained(model_config["checkpoint_path"]) model.eval() print("Fine-tuned MBart50 loaded successfully!") return model.to(DEVICE), tokenizer #type: ignore except Exception as e: raise RuntimeError(f"Failed to load fine-tuned MBart50 model: {str(e)}") @staticmethod def load_original_mbart50() -> Tuple[MBartForConditionalGeneration, MBart50Tokenizer]: """Load original MBart50 model and tokenizer.""" try: model_name = "facebook/mbart-large-50-many-to-many-mmt" model = MBartForConditionalGeneration.from_pretrained(model_name) tokenizer = MBart50Tokenizer.from_pretrained(model_name) model.eval() print("Original MBart50 loaded successfully!") return model.to(DEVICE), tokenizer #type: ignore except Exception as e: raise RuntimeError(f"Failed to load original MBart50 model: {str(e)}") @staticmethod def load_mt5() -> Tuple[MT5ForConditionalGeneration, MT5TokenizerFast]: """Load fine-tuned MT5 model and tokenizer.""" try: model_config = CONFIG["mt5"]["paths"] model = MT5ForConditionalGeneration.from_pretrained(model_config["base_model_name"]) peft_config = PeftConfig.from_pretrained(model_config["checkpoint_path"]) model = PeftModel.from_pretrained(model, model_config["checkpoint_path"]) tokenizer = MT5TokenizerFast.from_pretrained(model_config["checkpoint_path"]) model.eval() print("Fine-tuned MT5 loaded successfully!") return model.to(DEVICE), tokenizer #type: ignore except Exception as e: raise RuntimeError(f"Failed to load fine-tuned MT5 model: {str(e)}") @staticmethod def load_original_mt5() -> Tuple[MT5ForConditionalGeneration, MT5TokenizerFast]: """Load original MT5 model and tokenizer.""" try: model_name = "google/mt5-base" model = MT5ForConditionalGeneration.from_pretrained(model_name) tokenizer = MT5TokenizerFast.from_pretrained(model_name) model.eval() print("Original MT5 loaded successfully!") return model.to(DEVICE), tokenizer #type: ignore except Exception as e: raise RuntimeError(f"Failed to load original MT5 model: {str(e)}") class Translator: """Handles translation using different models.""" @staticmethod def translate_rbmt(text: str) -> str: """Translate using Rule-Based Machine Translation.""" try: translator = TransferBasedMT() return translator.translate(text) except Exception as e: raise RuntimeError(f"RBMT translation failed: {str(e)}") @staticmethod def translate_smt(text: str, smt) -> str: """Translate using Statistical Machine Translation.""" try: # return smt.translate_sentence(text) translation = smt.infer(text) return translation except Exception as e: raise RuntimeError(f"SMT translation failed: {str(e)}") @staticmethod def translate_mbart50( model: MBartForConditionalGeneration, tokenizer: MBart50Tokenizer, text: str ) -> str: """Translate using MBart50 model (fine-tuned or original).""" try: model_config = CONFIG["mbart50"]["args"] tokenizer.src_lang = model_config["src_lang"] inputs = tokenizer(text, return_tensors="pt", padding=True) inputs = {key: value.to(DEVICE) for key, value in inputs.items()} forced_bos_token_id = tokenizer.lang_code_to_id[model_config["tgt_lang"]] translated_tokens = model.generate( input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], forced_bos_token_id=forced_bos_token_id, max_length=128, num_beams=5, ) return tokenizer.decode(translated_tokens[0], skip_special_tokens=True) except Exception as e: raise RuntimeError(f"MBart50 translation failed: {str(e)}") @staticmethod def translate_mt5( model: MT5ForConditionalGeneration, tokenizer: MT5TokenizerFast, text: str ) -> str: """Translate using MT5 model (fine-tuned or original).""" try: prefix = CONFIG["mt5"]["args"]["prefix"] text = prefix + text inputs = tokenizer(text, return_tensors="pt", padding=True) inputs = {key: value.to(DEVICE) for key, value in inputs.items()} translated_tokens = model.generate( input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], max_length=128, num_beams=5, ) return tokenizer.decode(translated_tokens[0], skip_special_tokens=True) except Exception as e: raise RuntimeError(f"MT5 translation failed: {str(e)}") class Evaluator: """Handles evaluation of translation models.""" @staticmethod def load_test_data(test_file: str) -> List[Dict[str, str]]: """Load test data from CSV file.""" try: df = pd.read_csv(test_file) df = df[:100000] # Limit to 100,000 rows return [{"source": row["en"], "reference": row["vi"]} for _, row in df.iterrows()] except Exception as e: raise RuntimeError(f"Failed to load test data: {str(e)}") @staticmethod def compute_metrics(hypotheses: List[str], references: List[str], sources: List[str]) -> Dict[str, float]: """Compute translation evaluation metrics.""" try: metrics = {} bleu_metric = evaluate.load("sacrebleu") meteor_metric = evaluate.load("meteor") rouge_metric = evaluate.load("rouge") comet_metric = evaluate.load("comet") bertscore_metric = evaluate.load("bertscore") # BLEU metrics["SacreBLEU"] = bleu_metric.compute(predictions=hypotheses, references=references)["score"] / 100 #type: ignore # METEOR metrics["METEOR"] = meteor_metric.compute(predictions=hypotheses, references=references)["meteor"] #type: ignore # ROUGE rouge_results = rouge_metric.compute( predictions=hypotheses, references=references, rouge_types=["rouge1", "rouge2", "rougeL"], use_stemmer=True ) metrics["ROUGE-1"] = rouge_results["rouge1"] #type: ignore metrics["ROUGE-2"] = rouge_results["rouge2"] #type: ignore metrics["ROUGE-L"] = rouge_results["rougeL"] #type: ignore # BERTScore bertscore_results = bertscore_metric.compute( predictions=hypotheses, references=references, model_type="bert-base-multilingual-cased", lang="vi" ) metrics["BERTScore"] = sum(bertscore_results["f1"]) / len(bertscore_results["f1"]) #type: ignore # COMET comet_results = comet_metric.compute(predictions=hypotheses, references=references, sources=sources) metrics["COMET"] = sum(comet_results["scores"]) / len(comet_results["scores"]) #type: ignore return metrics except Exception as e: raise RuntimeError(f"Failed to compute metrics: {str(e)}") @staticmethod def evaluate_model( model_type: str, test_data: List[Dict[str, str]] ) -> Tuple[List[str], List[str], Dict[str, float]]: """Evaluate a translation model on test data.""" hypotheses, references, sources = [], [], [] try: if model_type == "rbmt": for item in tqdm(test_data, desc="Translating with RBMT"): translation = Translator.translate_rbmt(item["source"]) hypotheses.append(translation) references.append(item["reference"]) sources.append(item["source"]) elif model_type == "smt": for item in tqdm(test_data, desc="Translating with SMT"): smt = ModelLoader.load_smt() translation = Translator.translate_smt(item["source"], smt) hypotheses.append(translation) references.append(item["reference"]) sources.append(item["source"]) elif model_type == "mbart50": model, tokenizer = ModelLoader.load_mbart50() for item in tqdm(test_data, desc="Translating with fine-tuned mBART50"): translation = Translator.translate_mbart50(model, tokenizer, item["source"]) hypotheses.append(translation) references.append(item["reference"]) sources.append(item["source"]) elif model_type == "original_mbart50": model, tokenizer = ModelLoader.load_original_mbart50() for item in tqdm(test_data, desc="Translating with original mBART50"): translation = Translator.translate_mbart50(model, tokenizer, item["source"]) hypotheses.append(translation) references.append(item["reference"]) sources.append(item["source"]) elif model_type == "mt5": model, tokenizer = ModelLoader.load_mt5() for item in tqdm(test_data, desc="Translating with fine-tuned MT5"): translation = Translator.translate_mt5(model, tokenizer, item["source"]) hypotheses.append(translation) references.append(item["reference"]) sources.append(item["source"]) elif model_type == "original_mt5": model, tokenizer = ModelLoader.load_original_mt5() for item in tqdm(test_data, desc="Translating with original MT5"): translation = Translator.translate_mt5(model, tokenizer, item["source"]) hypotheses.append(translation) references.append(item["reference"]) sources.append(item["source"]) return hypotheses, references, Evaluator.compute_metrics(hypotheses, references, sources) if hypotheses else {} except Exception as e: raise RuntimeError(f"Evaluation failed for {model_type}: {str(e)}") def main(): """Main function to run model evaluation.""" args = parse_arguments() Path(args.output_dir).mkdir(parents=True, exist_ok=True) try: test_data = Evaluator.load_test_data(args.test_file) model_types = ["rbmt", "smt" "mbart50", "original_mbart50", "mt5", "original_mt5"] all_results = {} for model_type in model_types: print(f"\nEvaluating {model_type}...") hypotheses, references, metrics = Evaluator.evaluate_model(model_type, test_data) if metrics: all_results[model_type] = metrics print(f"Metrics for {model_type}:") for metric, value in metrics.items(): print(f"{metric}: {value:.4f}") # Save translations translations = [ {"source": item["source"], "reference": ref, "hypothesis": hyp} for item, ref, hyp in zip(test_data, references, hypotheses) ] with open( Path(args.output_dir) / f"{model_type}_translations.json", "w", encoding="utf-8" ) as f: json.dump(translations, f, ensure_ascii=False, indent=2) # Save all metrics with open(Path(args.output_dir) / "metrics.json", "w", encoding="utf-8") as f: json.dump(all_results, f, indent=2) except Exception as e: print(f"Error: {str(e)}", file=sys.stderr) sys.exit(1) if __name__ == "__main__": main()