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())