translator_engine.py

import torch
import time
import re
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, Trainer, TrainingArguments
from datasets import Dataset 
import pandas as pd
from config import MODEL_MAP, BATCH_SIZE_LIMIT, DATASET_PATH, FINETUNED_MODEL_BASE_DIR, DEFAULT_MODEL_CHOICE
from utils import preprocess_text, load_iwslt_dataset, preprocess_pair
import os 
import json 

# --- THÊM CÁC IMPORTS MỚI CHO ĐÁNH GIÁ METRIC ---
import evaluate
import numpy as np

# Tải metric BLEU một lần
metric_bleu = evaluate.load("sacrebleu") 

class TranslationEngine:
    def __init__(self):
        self.translators = {}
        self.tokenizers = {}
        self.models = {}
        self.train_data, self.val_data, self.test_data = self.load_dataset()

    def load_dataset(self):
        """Load IWSLT15 dataset for training or evaluation"""
        train_data, val_data, test_data = load_iwslt_dataset(DATASET_PATH)
        if train_data is None:
            print("Failed to load dataset. Using default behavior without dataset.")
            return None, None, None
        return train_data, val_data, test_data

    def load_model_and_tokenizer(self, model_choice_display_name): # Sử dụng tên tham số nhất quán
        # Kiểm tra xem mô hình này đã được tải vào bộ nhớ chưa
        if model_choice_display_name in self.translators:
            return self.translators[model_choice_display_name], self.tokenizers[model_choice_display_name]

        print(f"Loading model: {model_choice_display_name}")
        
        # Lấy thông tin model từ MODEL_MAP
        model_info = MODEL_MAP.get(model_choice_display_name)
        if not model_info:
            print(f"  !!! ERROR: Model '{model_choice_display_name}' not found in MODEL_MAP. Falling back to default.")
            model_info = MODEL_MAP[DEFAULT_MODEL_CHOICE] # Fallback nếu tên không hợp lệ
        
        model_type = model_info["type"]
        model_identifier = model_info["identifier"] # Có thể là ID HF hoặc tên thư mục cục bộ

        tokenizer = None
        model = None
        
        if model_type == "hub":
            # Tải mô hình từ Hugging Face Hub (pretrained)
            print(f"  --> Loading pre-trained model from Hugging Face Hub: {model_identifier}")
            tokenizer = AutoTokenizer.from_pretrained(model_identifier)
            model = AutoModelForSeq2SeqLM.from_pretrained(model_identifier)
        elif model_type == "local_finetuned":
            # Tải mô hình đã fine-tuned từ local
            finetuned_model_path_local = os.path.join(FINETUNED_MODEL_BASE_DIR, model_identifier)
            
            print(f"  --> Loading fine-tuned model from local directory: {finetuned_model_path_local}")
            
            # --- Gỡ lỗi chi tiết các file quan trọng ---
            config_json_path = os.path.join(finetuned_model_path_local, "config.json")
            tokenizer_config_json_path = os.path.join(finetuned_model_path_local, "tokenizer_config.json")
            model_safetensors_path = os.path.join(finetuned_model_path_local, "model.safetensors")
            
            print(f"  - config.json exists: {os.path.exists(config_json_path)}")
            print(f"  - tokenizer_config.json exists: {os.path.exists(tokenizer_config_json_path)}")
            print(f"  - model.safetensors exists: {os.path.exists(model_safetensors_path)}")
            print(f"  - source.spm exists: {os.path.exists(os.path.join(finetuned_model_path_local, 'source.spm'))}")
            print(f"  - target.spm exists: {os.path.exists(os.path.join(finetuned_model_path_local, 'target.spm'))}")

            try:
                # Cố gắng tải CẢ tokenizer VÀ model từ thư mục fine-tuned cục bộ
                tokenizer = AutoTokenizer.from_pretrained(finetuned_model_path_local)
                model = AutoModelForSeq2SeqLM.from_pretrained(finetuned_model_path_local)
                print(f"  --> Successfully loaded fine-tuned model and tokenizer from local path.")
            except Exception as e:
                print(f"  !!! ERROR: Failed to load fine-tuned model/tokenizer from '{finetuned_model_path_local}': {e}")
                print(f"  --> Falling back to loading original pre-trained model: {MODEL_MAP[DEFAULT_MODEL_CHOICE]['identifier']}")
                # Nếu tải fine-tuned thất bại, tải lại mô hình pretrained gốc
                tokenizer = AutoTokenizer.from_pretrained(MODEL_MAP[DEFAULT_MODEL_CHOICE]['identifier'])
                model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_MAP[DEFAULT_MODEL_CHOICE]['identifier'])
        else:
            # Nếu thư mục fine-tuned model không tồn tại
            print(f"  --> Local fine-tuned directory '{finetuned_model_path_local}' not found. Falling back to pre-trained model: {MODEL_MAP[DEFAULT_MODEL_CHOICE]['identifier']}")
            tokenizer = AutoTokenizer.from_pretrained(MODEL_MAP[DEFAULT_MODEL_CHOICE]['identifier'])
            model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_MAP[DEFAULT_MODEL_CHOICE]['identifier'])
        
        # --- KẾT THÚC LOGIC TẢI MÔ HÌNH ---
        
        device = 0 if torch.cuda.is_available() else -1
        if device == 0:
            model = model.cuda()
            print("Using GPU for faster inference")
        
        translator = pipeline(
            "translation",
            model=model,
            tokenizer=tokenizer,
            device=device,
            batch_size=BATCH_SIZE_LIMIT,
            torch_dtype=torch.float16 if device == 0 and torch.cuda.is_available() else None
        )
        
        self.translators[model_choice_display_name] = translator
        self.tokenizers[model_choice_display_name] = tokenizer
        self.models[model_choice_display_name] = model
        
        print(f"Model {model_choice_display_name} loaded successfully")
        
        return self.translators[model_choice_display_name], self.tokenizers[model_choice_display_name]

    def fine_tune_model(self, model_choice_display_name, output_dir_base=FINETUNED_MODEL_BASE_DIR): # Tham số này giờ là display name
        if not self.train_data or not self.val_data:
            print("Dataset not loaded or missing validation split. Cannot fine-tune model.")
            return

        # Xác định ID mô hình gốc (từ display name của bản fine-tuned) để tạo thư mục lưu
        # Chúng ta giả định rằng model_choice_display_name sẽ là "English to Vietnamese (OpusMT) - Fine-tuned"
        # và identifier của nó sẽ là "English to Vietnamese (OpusMT)" như trong config.py
        if model_choice_display_name == "English to Vietnamese (OpusMT) - Fine-tuned":
            model_identifier_for_save = MODEL_MAP[model_choice_display_name]["identifier"]
        else:
            # Fallback nếu người dùng chọn model khác để fine-tune (ví dụ: bản pretrained)
            # Trong trường hợp này, chúng ta sẽ lưu nó theo identifier của nó.
            model_identifier_for_save = MODEL_MAP[model_choice_display_name]["identifier"]

        # Tải mô hình và tokenizer ban đầu (sẽ được dùng để fine-tune)
        # load_model_and_tokenizer sẽ tải mô hình gốc (pretrained) khi bắt đầu fine-tune
        # vì chúng ta muốn fine-tune từ một điểm khởi đầu sạch.
        # Sử dụng tên hiển thị của bản pretrained để tải đúng model_id gốc
        pretrained_display_name = "English to Vietnamese (OpusMT) - Pretrained"
        translator, tokenizer = self.load_model_and_tokenizer(pretrained_display_name) 
        model = self.models[pretrained_display_name] # Lấy instance mô hình đã được tải

        # Đường dẫn lưu model đã fine-tuned cụ thể
        current_finetune_output_dir = os.path.join(output_dir_base, model_identifier_for_save)
        
        if not os.path.exists(current_finetune_output_dir):
            os.makedirs(current_finetune_output_dir)
            print(f"Đã tạo thư mục lưu mô hình fine-tuned: {current_finetune_output_dir}")

        def tokenize_data(example):
            en_texts_batch = []
            vi_texts_batch = []
            
            for pair_dict in example["translation"]:
                en_text, vi_text = preprocess_pair(pair_dict["en"], pair_dict["vi"])
                en_texts_batch.append(en_text)
                vi_texts_batch.append(vi_text)
            
            inputs = tokenizer(
                en_texts_batch,
                truncation=True,
                padding="max_length",
                max_length=128, 
                return_tensors="pt"
            )
            labels = tokenizer(
                vi_texts_batch,
                truncation=True,
                padding="max_length",
                max_length=128, 
                return_tensors="pt"
            )
            inputs["labels"] = labels["input_ids"]
            return inputs

        print("Tokenizing training data...")
        tokenized_train = self.train_data.map(tokenize_data, batched=True, remove_columns=self.train_data.column_names)
        print("Tokenizing validation data...")
        tokenized_val = self.val_data.map(tokenize_data, batched=True, remove_columns=self.val_data.column_names)

        def compute_metrics(eval_preds):
            preds, labels = eval_preds
            if isinstance(preds, tuple):
                preds = preds[0]
            
            decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
            decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

            decoded_preds = [pred.strip() for pred in decoded_preds]
            decoded_labels = [[label.strip()] for label in decoded_labels] 

            result = metric_bleu.compute(predictions=decoded_preds, references=decoded_labels)
            
            result = {"bleu": round(result["score"], 2)}
            return result

        training_args = TrainingArguments(
            output_dir=current_finetune_output_dir,
            evaluation_strategy="epoch",
            learning_rate=2e-5,
            per_device_train_batch_size=8,
            per_device_eval_batch_size=8,
            num_train_epochs=1, 
            weight_decay=0.01,
            save_strategy="epoch",
            load_best_model_at_end=True,
            logging_dir=f'{current_finetune_output_dir}/logs',
            logging_steps=500,
            save_total_limit=1,
            report_to="none", 
            predict_with_generate=True,
        )

        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=tokenized_train,
            eval_dataset=tokenized_val,
            tokenizer=tokenizer,          
            compute_metrics=compute_metrics, 
        )

        print("Starting fine-tuning...")
        trainer.train()
        print(f"Fine-tuning completed. Model saved to {current_finetune_output_dir}")

        trainer.save_model(current_finetune_output_dir)
        tokenizer.save_pretrained(current_finetune_output_dir)
        print(f"Mô hình và tokenizer đã được lưu đầy đủ vào: {current_finetune_output_dir}")

        # Sau khi fine-tune xong, cập nhật pipeline cho bản fine-tuned để sử dụng ngay lập tức
        print(f"Reloading fine-tuned model into pipeline from {current_finetune_output_dir} for immediate inference...")
        try:
            reloaded_tokenizer = AutoTokenizer.from_pretrained(current_finetune_output_dir)
            reloaded_model = AutoModelForSeq2SeqLM.from_pretrained(current_finetune_output_dir)
            if torch.cuda.is_available():
                reloaded_model = reloaded_model.cuda()
            
            # CẬP NHẬT TRÌNH DỊCH VỚI TÊN HIỂN THỊ CỦA BẢN FINE-TUNED (model_choice_display_name)
            self.translators[model_choice_display_name] = pipeline(
                "translation",
                model=reloaded_model, 
                tokenizer=reloaded_tokenizer,
                device=0 if torch.cuda.is_available() else -1,
                batch_size=BATCH_SIZE_LIMIT,
                torch_dtype=torch.float16 if torch.cuda.is_available() else None
            )
            self.models[model_choice_display_name] = reloaded_model 
            self.tokenizers[model_choice_display_name] = reloaded_tokenizer 
            print(f"Fine-tuned model for '{model_choice_display_name}' reloaded successfully for inference.")
        except Exception as e_reload:
            print(f"!!! CRITICAL: Failed to reload saved model for immediate inference: {e_reload}")
            # Nếu không thể tải lại bản fine-tuned, hãy đảm bảo tải bản pretrained để app không crash
            print(f"!!! Reverting to pre-trained model for current inference.")
            pretrained_info = MODEL_MAP["English to Vietnamese (OpusMT) - Pretrained"]
            tokenizer = AutoTokenizer.from_pretrained(pretrained_info["identifier"])
            model = AutoModelForSeq2SeqLM.from_pretrained(pretrained_info["identifier"])
            if torch.cuda.is_available():
                model = model.cuda()
            self.translators[model_choice_display_name] = pipeline(
                "translation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1,
                batch_size=BATCH_SIZE_LIMIT, torch_dtype=torch.float16 if torch.cuda.is_available() else None
            )
            self.models[model_choice_display_name] = model
            self.tokenizers[model_choice_display_name] = tokenizer


    def clean_vietnamese_output(self, text):
        if not text:
            return ""
        text = text.replace("_", " ")
        text = re.sub(r"\s+", " ", text).strip()
        return text

    def translate_batch(self, texts, translator, max_length, num_beams=1):
        try:
            results = translator(
                texts,
                max_length=max_length,
                num_beams=num_beams,
                batch_size=BATCH_SIZE_LIMIT
            )
            cleaned_results = [self.clean_vietnamese_output(r["translation_text"]) for r in results]
            return cleaned_results
        except Exception as e:
            print(f"Batch translation error: {e}")
            results = []
            for text in texts:
                try:
                    result = translator(
                        text,
                        max_length=max_length,
                        num_beams=num_beams
                    )[0]["translation_text"]
                    results.append(self.clean_vietnamese_output(result)) 
                except Exception as e2:
                    print(f"Individual translation error for '{text}': {e2}")
                    results.append(f"Translation error: {str(e2)}")
            return results

    def translate_single_text(self, text, model_choice_display_name, max_length=512, num_beams=1): # Đổi tên tham số
        if not text:
            return None
        
        start_time_total = time.time()
        processed_text = preprocess_text(text) 
        
        # Gọi load_model_and_tokenizer với tên hiển thị
        translator, tokenizer = self.load_model_and_tokenizer(model_choice_display_name) 
        
        tokens = tokenizer(processed_text, return_tensors="pt", truncation=False)
        if len(tokens["input_ids"][0]) > max_length:
            sentences = [s.strip() + '.' for s in processed_text.split('.') if s.strip()]
            if not sentences and processed_text.strip():
                 sentences = [processed_text.strip()]
            
            results = self.translate_batch(sentences, translator, max_length, num_beams)
            result = ' '.join(results)
        else:
            results = self.translate_batch([processed_text], translator, max_length, num_beams)
            result = results[0]
        
        result = self.clean_vietnamese_output(result)
        
        print(f"Translation completed in: {time.time() - start_time_total:.2f}s")
        return result

    def translate_file_content(self, texts, model_choice_display_name, max_length=512, num_beams=1): # Đổi tên tham số
        start_time = time.time()
        
        processed_texts = [preprocess_text(t) for t in texts]
        translator, _ = self.load_model_and_tokenizer(model_choice_display_name) # Gọi với tên hiển thị
        
        translated_results_raw = self.translate_batch(processed_texts, translator, max_length, num_beams)
        
        translated_results_cleaned = [self.clean_vietnamese_output(r) for r in translated_results_raw]

        print(f"File translation completed in: {time.time() - start_time:.2f}s")
        return list(zip(processed_texts, translated_results_cleaned))

    def warm_up_model(self, model_choice_display_name): # Đổi tên tham số
        try:
            translator, _ = self.load_model_and_tokenizer(model_choice_display_name) # Gọi với tên hiển thị
            _ = translator("Hello", max_length=50, num_beams=1)
            print(f"Model {model_choice_display_name} warmed up successfully")
        except Exception as e:
            print(f"Warm-up failed for {model_choice_display_name}: {e}")

    def get_loaded_models(self):
        return list(self.translators.keys())