Upload model source code

src/1_build_shared_vocab.py

@@ -0,0 +1,214 @@
+"""
+BƯỚC 1: TẠO BỘ TỪ ĐIỂN CHUNG (SHARED VOCABULARY)
+Train BPE tokenizer từ text gốc để tạo shared vocabulary cho cả tiếng Việt và tiếng Anh
+"""
+
+from pathlib import Path
+from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers
+
+# ============================================================================
+# PATH CONFIGURATION
+# ============================================================================
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DATA_DIR = PROJECT_ROOT / 'data'
+RAW_DATA_DIR = DATA_DIR / 'raw'
+PROCESSED_DATA_DIR = DATA_DIR / 'processed'
+
+# ============================================================================
+# BUILD SHARED VOCABULARY
+# ============================================================================
+
+def build_shared_vocab(
+    vocab_size=32000,
+    min_frequency=2,
+    special_tokens=None,
+    output_path=None
+):
+    """
+    Train BPE tokenizer từ text files để tạo shared vocabulary
+    
+    Args:
+        vocab_size: Kích thước vocabulary (mặc định 32000)
+        min_frequency: Tần suất tối thiểu để giữ lại token
+        special_tokens: List special tokens [PAD, SOS, EOS, UNK]
+        output_path: Đường dẫn lưu tokenizer
+        
+    Returns:
+        tokenizer: Trained tokenizer
+    """
+    if special_tokens is None:
+        # Quan trọng: Đặt Special Tokens ĐÚNG THỨ TỰ
+        # 0=PAD, 1=SOS, 2=EOS, 3=UNK (phải khớp với code hiện tại)
+        special_tokens = ["<pad>", "<sos>", "<eos>", "<unk>"]
+    
+    if output_path is None:
+        output_path = PROCESSED_DATA_DIR / 'tokenizer_shared.json'
+    
+    print("="*70)
+    print("TẠO SHARED VOCABULARY (BPE TOKENIZER)")
+    print("="*70)
+    
+    # 1. Cấu hình Tokenizer BPE
+    print("\n1. Khởi tạo BPE Tokenizer...")
+    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
+    
+    # Pre-tokenizer: ByteLevel (xử lý Unicode tốt)
+    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+    tokenizer.decoder = decoders.ByteLevel()
+    
+    # 2. Trainer configuration
+    print(f"2. Cấu hình trainer:")
+    print(f"   - Vocab size: {vocab_size}")
+    print(f"   - Min frequency: {min_frequency}")
+    print(f"   - Special tokens: {special_tokens}")
+    
+    trainer = trainers.BpeTrainer(
+        vocab_size=vocab_size,
+        special_tokens=special_tokens,
+        min_frequency=min_frequency,
+        show_progress=True
+    )
+    
+    # 3. Tìm file text gốc
+    # Ưu tiên: MTet CSV (nếu có MTET_MAX_ROWS) > Raw files > cleaned_data.json
+    files = []
+    
+    # Option 0: Dùng MTet CSV nếu có MTET_MAX_ROWS set
+    from data_preprocessing import CUSTOM_CSV_PATH, CUSTOM_CSV_MAX_ROWS
+    if CUSTOM_CSV_PATH.exists() and CUSTOM_CSV_MAX_ROWS and CUSTOM_CSV_MAX_ROWS > 0:
+        print(f"\n3. Ưu tiên dùng MTet CSV với {CUSTOM_CSV_MAX_ROWS:,} câu:")
+        print(f"   - {CUSTOM_CSV_PATH}")
+        print(f"   Tạo temporary text files từ CSV...")
+        
+        import json
+        import tempfile
+        import pandas as pd
+        
+        # Đọc CSV
+        df = pd.read_csv(CUSTOM_CSV_PATH, nrows=CUSTOM_CSV_MAX_ROWS)
+        if {'src', 'tgt'}.issubset(df.columns):
+            df = df.rename(columns={'src': 'vi', 'tgt': 'en'})
+        elif {'vi', 'en'}.issubset(df.columns):
+            pass
+        else:
+            raise ValueError("CSV cần có cột 'src'/'tgt' hoặc 'vi'/'en'")
+        
+        df = df[['vi', 'en']].dropna().reset_index(drop=True)
+        
+        # Tạo temp files
+        temp_dir = Path(tempfile.mkdtemp())
+        temp_vi = temp_dir / 'train.vi'
+        temp_en = temp_dir / 'train.en'
+        
+        with open(temp_vi, 'w', encoding='utf-8') as f_vi, \
+             open(temp_en, 'w', encoding='utf-8') as f_en:
+            for _, row in df.iterrows():
+                f_vi.write(str(row['vi']).strip() + '\n')
+                f_en.write(str(row['en']).strip() + '\n')
+        
+        files = [str(temp_vi), str(temp_en)]
+        print(f"   ✓ Đã tạo temp files: {len(df):,} cặp câu từ MTet CSV")
+    
+    # Option 1: Dùng raw files
+    elif (RAW_DATA_DIR / 'train.vi').exists() and (RAW_DATA_DIR / 'train.en').exists():
+        train_vi = RAW_DATA_DIR / 'train.vi'
+        train_en = RAW_DATA_DIR / 'train.en'
+        files = [str(train_vi), str(train_en)]
+        print(f"\n3. Tìm thấy raw files:")
+        print(f"   - {train_vi}")
+        print(f"   - {train_en}")
+    
+    # Option 2: Fallback - dùng cleaned_data.json
+    else:
+        # Dùng cleaned_data.json nếu có
+        cleaned_data_path = PROCESSED_DATA_DIR / 'cleaned_data.json'
+        if cleaned_data_path.exists():
+            print(f"\n3. Không tìm thấy raw files, dùng cleaned_data.json")
+            print(f"   Tạo temporary text files từ cleaned_data.json...")
+            
+            import json
+            import tempfile
+            
+            with open(cleaned_data_path, 'r', encoding='utf-8') as f:
+                cleaned_data = json.load(f)
+            
+            # Tạo temp files
+            temp_dir = Path(tempfile.mkdtemp())
+            temp_vi = temp_dir / 'train.vi'
+            temp_en = temp_dir / 'train.en'
+            
+            with open(temp_vi, 'w', encoding='utf-8') as f_vi, \
+                 open(temp_en, 'w', encoding='utf-8') as f_en:
+                for item in cleaned_data.get('train', []):
+                    f_vi.write(item['vi'] + '\n')
+                    f_en.write(item['en'] + '\n')
+            
+            files = [str(temp_vi), str(temp_en)]
+            print(f"   ✓ Đã tạo temp files: {len(cleaned_data.get('train', []))} cặp câu")
+        else:
+            raise FileNotFoundError(
+                f"Không tìm thấy dữ liệu để train tokenizer!\n"
+                f"Vui lòng đảm bảo có một trong các file sau:\n"
+                f"  - {train_vi} và {train_en}\n"
+                f"  - {cleaned_data_path}"
+            )
+    
+    # 4. Train tokenizer
+    print(f"\n4. Đang train tokenizer từ {len(files)} file(s)...")
+    print(f"   (Quá trình này có thể mất 5-10 phút tùy kích thước dữ liệu)")
+    
+    tokenizer.train(files, trainer)
+    
+    # 5. Lưu tokenizer
+    PROCESSED_DATA_DIR.mkdir(parents=True, exist_ok=True)
+    tokenizer.save(str(output_path))
+    
+    print(f"\n5. ✓ Đã lưu tokenizer: {output_path}")
+    
+    # 6. Kiểm tra special tokens
+    print(f"\n6. Kiểm tra Special Tokens:")
+    print(f"   - PAD (<pad>): ID {tokenizer.token_to_id('<pad>')} (phải là 0)")
+    print(f"   - SOS (<sos>): ID {tokenizer.token_to_id('<sos>')} (phải là 1)")
+    print(f"   - EOS (<eos>): ID {tokenizer.token_to_id('<eos>')} (phải là 2)")
+    print(f"   - UNK (<unk>): ID {tokenizer.token_to_id('<unk>')} (phải là 3)")
+    
+    # 7. Thống kê
+    vocab_size_actual = tokenizer.get_vocab_size()
+    print(f"\n7. Thống kê:")
+    print(f"   - Vocab size thực tế: {vocab_size_actual}")
+    print(f"   - File output: {output_path}")
+    
+    print("\n" + "="*70)
+    print("✓ HOÀN TẤT TẠO SHARED VOCABULARY!")
+    print("="*70)
+    
+    return tokenizer
+
+# ============================================================================
+# MAIN
+# ============================================================================
+
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser(description='Build Shared Vocabulary')
+    parser.add_argument('--vocab_size', type=int, default=32000,
+                       help='Vocabulary size (default: 32000)')
+    parser.add_argument('--min_frequency', type=int, default=2,
+                       help='Minimum token frequency (default: 2)')
+    parser.add_argument('--output', type=str, default=None,
+                       help='Output path for tokenizer (default: data/processed/tokenizer_shared.json)')
+    
+    args = parser.parse_args()
+    
+    tokenizer = build_shared_vocab(
+        vocab_size=args.vocab_size,
+        min_frequency=args.min_frequency,
+        output_path=args.output
+    )
+    
+    print("\n📝 BƯỚC TIẾP THEO:")
+    print("   Chạy: python src/2_encode_data.py")
+    print("   để encode lại dữ liệu với shared vocabulary này")
+

src/2_encode_data.py

@@ -0,0 +1,335 @@
+"""
+BƯỚC 2: ENCODE LẠI DỮ LIỆU VỚI SHARED VOCABULARY
+Dùng tokenizer_shared.json để encode lại tất cả dữ liệu từ text gốc
+"""
+
+import json
+import pickle
+from pathlib import Path
+from tokenizers import Tokenizer
+from tqdm import tqdm
+
+# ============================================================================
+# PATH CONFIGURATION
+# ============================================================================
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DATA_DIR = PROJECT_ROOT / 'data'
+RAW_DATA_DIR = DATA_DIR / 'raw'
+PROCESSED_DATA_DIR = DATA_DIR / 'processed'
+
+# ============================================================================
+# ENCODE DATA WITH SHARED VOCABULARY
+# ============================================================================
+
+def encode_file_with_shared_vocab(src_file, tgt_file, tokenizer, split_name='train'):
+    """
+    Encode một cặp file (source và target) với shared tokenizer
+    
+    Args:
+        src_file: Path to source file (Vietnamese)
+        tgt_file: Path to target file (English)
+        tokenizer: Shared tokenizer
+        split_name: Name of split (train/validation/test)
+        
+    Returns:
+        data: List of (src_ids, tgt_ids) tuples
+    """
+    print(f"\nĐang encode {split_name} set...")
+    
+    # Lấy special token IDs
+    SOS_ID = tokenizer.token_to_id("<sos>")  # Thường là 1
+    EOS_ID = tokenizer.token_to_id("<eos>")  # Thường là 2
+    
+    if SOS_ID is None or EOS_ID is None:
+        raise ValueError("Không tìm thấy <sos> hoặc <eos> trong tokenizer!")
+    
+    data = []
+    
+    # Đọc và encode từng cặp câu
+    with open(src_file, 'r', encoding='utf-8') as f_src, \
+         open(tgt_file, 'r', encoding='utf-8') as f_tgt:
+        
+        src_lines = f_src.readlines()
+        tgt_lines = f_tgt.readlines()
+        
+        if len(src_lines) != len(tgt_lines):
+            print(f"⚠️  Warning: Số dòng không khớp ({len(src_lines)} vs {len(tgt_lines)})")
+            min_len = min(len(src_lines), len(tgt_lines))
+            src_lines = src_lines[:min_len]
+            tgt_lines = tgt_lines[:min_len]
+        
+        # Encode từng cặp
+        for src_text, tgt_text in tqdm(zip(src_lines, tgt_lines), 
+                                       total=len(src_lines), 
+                                       desc=f"Encoding {split_name}"):
+            # Clean và encode
+            src_text = src_text.strip()
+            tgt_text = tgt_text.strip()
+            
+            if not src_text or not tgt_text:
+                continue
+            
+            # Encode text -> IDs
+            src_encoded = tokenizer.encode(src_text)
+            tgt_encoded = tokenizer.encode(tgt_text)
+            
+            src_ids = src_encoded.ids
+            tgt_ids = tgt_encoded.ids
+            
+            # Thêm SOS và EOS thủ công để đảm bảo đúng format
+            # Format: [SOS, ...token_ids..., EOS]
+            src_full = [SOS_ID] + src_ids + [EOS_ID]
+            tgt_full = [SOS_ID] + tgt_ids + [EOS_ID]
+            
+            # Filter quá dài (giới hạn 100 tokens)
+            if len(src_full) > 100:
+                src_full = src_full[:99] + [EOS_ID]
+            if len(tgt_full) > 100:
+                tgt_full = tgt_full[:99] + [EOS_ID]
+            
+            data.append((src_full, tgt_full))
+    
+    print(f"  ✓ Đã encode {len(data)} cặp câu")
+    return data
+
+def encode_from_cleaned_data(tokenizer, cleaned_data_path):
+    """
+    Encode từ cleaned_data.json nếu không có raw files
+    
+    Args:
+        tokenizer: Shared tokenizer
+        cleaned_data_path: Path to cleaned_data.json
+        
+    Returns:
+        processed_data: Dict với keys 'train', 'validation', 'test'
+    """
+    print("\nĐang encode từ cleaned_data.json...")
+    
+    with open(cleaned_data_path, 'r', encoding='utf-8') as f:
+        cleaned_data = json.load(f)
+    
+    SOS_ID = tokenizer.token_to_id("<sos>")
+    EOS_ID = tokenizer.token_to_id("<eos>")
+    
+    processed_data = {
+        'train': [],
+        'validation': [],
+        'test': []
+    }
+    
+    for split in ['train', 'validation', 'test']:
+        print(f"\nĐang encode {split} set...")
+        
+        for item in tqdm(cleaned_data.get(split, []), desc=f"Encoding {split}"):
+            src_text = item['vi'].strip()
+            tgt_text = item['en'].strip()
+            
+            if not src_text or not tgt_text:
+                continue
+            
+            # Encode
+            src_encoded = tokenizer.encode(src_text)
+            tgt_encoded = tokenizer.encode(tgt_text)
+            
+            src_ids = [SOS_ID] + src_encoded.ids + [EOS_ID]
+            tgt_ids = [SOS_ID] + tgt_encoded.ids + [EOS_ID]
+            
+            # Filter quá dài
+            if len(src_ids) > 100:
+                src_ids = src_ids[:99] + [EOS_ID]
+            if len(tgt_ids) > 100:
+                tgt_ids = tgt_ids[:99] + [EOS_ID]
+            
+            processed_data[split].append((src_ids, tgt_ids))
+        
+        print(f"  ✓ {split}: {len(processed_data[split])} cặp câu")
+    
+    return processed_data
+
+def main():
+    """
+    Main function để encode lại toàn bộ dữ liệu
+    """
+    print("="*70)
+    print("ENCODE LẠI DỮ LIỆU VỚI SHARED VOCABULARY")
+    print("="*70)
+    
+    # 1. Load shared tokenizer
+    tokenizer_path = PROCESSED_DATA_DIR / 'tokenizer_shared.json'
+    
+    if not tokenizer_path.exists():
+        raise FileNotFoundError(
+            f"Không tìm thấy tokenizer_shared.json!\n"
+            f"Vui lòng chạy: python src/1_build_shared_vocab.py trước"
+        )
+    
+    print(f"\n1. Đang load tokenizer từ: {tokenizer_path}")
+    tokenizer = Tokenizer.from_file(str(tokenizer_path))
+    print(f"   ✓ Vocab size: {tokenizer.get_vocab_size()}")
+    
+    # 2. Kiểm tra và encode data
+    # Ưu tiên: MTet CSV (nếu có MTET_MAX_ROWS) > Raw files > cleaned_data.json
+    processed_data = {}
+    
+    # Option 0: Dùng MTet CSV nếu có MTET_MAX_ROWS set
+    from data_preprocessing import CUSTOM_CSV_PATH, CUSTOM_CSV_MAX_ROWS
+    if CUSTOM_CSV_PATH.exists() and CUSTOM_CSV_MAX_ROWS and CUSTOM_CSV_MAX_ROWS > 0:
+        print("\n2. Ưu tiên dùng MTet CSV với 1M câu...")
+        import pandas as pd
+        
+        # Đọc CSV
+        df = pd.read_csv(CUSTOM_CSV_PATH, nrows=CUSTOM_CSV_MAX_ROWS)
+        if {'src', 'tgt'}.issubset(df.columns):
+            df = df.rename(columns={'src': 'vi', 'tgt': 'en'})
+        elif {'vi', 'en'}.issubset(df.columns):
+            pass
+        else:
+            raise ValueError("CSV cần có cột 'src'/'tgt' hoặc 'vi'/'en'")
+        
+        df = df[['vi', 'en']].dropna().reset_index(drop=True)
+        total = len(df)
+        
+        # Split: 90% train, 5% val, 5% test
+        train_end = int(total * 0.90)
+        val_end = train_end + int(total * 0.05)
+        
+        train_df = df.iloc[:train_end]
+        val_df = df.iloc[train_end:val_end]
+        test_df = df.iloc[val_end:]
+        
+        print(f"   - Total: {total:,} cặp câu")
+        print(f"   - Train: {len(train_df):,} ({len(train_df)/total*100:.1f}%)")
+        print(f"   - Validation: {len(val_df):,} ({len(val_df)/total*100:.1f}%)")
+        print(f"   - Test: {len(test_df):,} ({len(test_df)/total*100:.1f}%)")
+        
+        # Encode từng split
+        SOS_ID = tokenizer.token_to_id("<sos>")
+        EOS_ID = tokenizer.token_to_id("<eos>")
+        
+        def encode_split(split_df, split_name):
+            data = []
+            for _, row in tqdm(split_df.iterrows(), total=len(split_df), desc=f"Encoding {split_name}"):
+                src_text = str(row['vi']).strip()
+                tgt_text = str(row['en']).strip()
+                
+                if not src_text or not tgt_text:
+                    continue
+                
+                src_encoded = tokenizer.encode(src_text)
+                tgt_encoded = tokenizer.encode(tgt_text)
+                
+                src_ids = [SOS_ID] + src_encoded.ids + [EOS_ID]
+                tgt_ids = [SOS_ID] + tgt_encoded.ids + [EOS_ID]
+                
+                if len(src_ids) > 100:
+                    src_ids = src_ids[:99] + [EOS_ID]
+                if len(tgt_ids) > 100:
+                    tgt_ids = tgt_ids[:99] + [EOS_ID]
+                
+                data.append((src_ids, tgt_ids))
+            return data
+        
+        processed_data['train'] = encode_split(train_df, 'train')
+        processed_data['validation'] = encode_split(val_df, 'validation')
+        processed_data['test'] = encode_split(test_df, 'test')
+    
+    # Option 1: Có raw files
+    elif all(f.exists() for f in [RAW_DATA_DIR / 'train.vi', RAW_DATA_DIR / 'train.en']):
+        train_vi = RAW_DATA_DIR / 'train.vi'
+        train_en = RAW_DATA_DIR / 'train.en'
+        dev_vi = RAW_DATA_DIR / 'dev.vi'
+        dev_en = RAW_DATA_DIR / 'dev.en'
+        test_vi = RAW_DATA_DIR / 'test.vi'
+        test_en = RAW_DATA_DIR / 'test.en'
+        print("\n2. Tìm thấy raw files, đang encode...")
+        
+        # Encode train
+        train_data = encode_file_with_shared_vocab(
+            train_vi, train_en, tokenizer, 'train'
+        )
+        processed_data['train'] = train_data
+        
+        # Encode validation (nếu có)
+        if dev_vi.exists() and dev_en.exists():
+            val_data = encode_file_with_shared_vocab(
+                dev_vi, dev_en, tokenizer, 'validation'
+            )
+            processed_data['validation'] = val_data
+        else:
+            print("\n⚠️  Không tìm thấy dev files, bỏ qua validation set")
+            processed_data['validation'] = []
+        
+        # Encode test (nếu có)
+        if test_vi.exists() and test_en.exists():
+            test_data = encode_file_with_shared_vocab(
+                test_vi, test_en, tokenizer, 'test'
+            )
+            processed_data['test'] = test_data
+        else:
+            print("\n⚠️  Không tìm thấy test files, bỏ qua test set")
+            processed_data['test'] = []
+    
+    # Option 2: Dùng cleaned_data.json
+    else:
+        cleaned_data_path = PROCESSED_DATA_DIR / 'cleaned_data.json'
+        
+        if not cleaned_data_path.exists():
+            raise FileNotFoundError(
+                f"Không tìm thấy raw files hoặc cleaned_data.json!\n"
+                f"Vui lòng đảm bảo có một trong các file sau:\n"
+                f"  - {train_vi}, {train_en}\n"
+                f"  - {cleaned_data_path}"
+            )
+        
+        print("\n2. Không tìm thấy raw files, dùng cleaned_data.json...")
+        processed_data = encode_from_cleaned_data(tokenizer, cleaned_data_path)
+    
+    # 3. Lưu processed data
+    print("\n3. Đang lưu processed data...")
+    
+    # Format: {'train': [(src_ids, tgt_ids), ...], ...}
+    output_path = PROCESSED_DATA_DIR / 'processed_data_shared.pkl'
+    
+    with open(output_path, 'wb') as f:
+        pickle.dump(processed_data, f)
+    
+    print(f"   ✓ Đã lưu: {output_path}")
+    
+    # 4. Thống kê
+    print("\n4. Thống kê:")
+    print(f"   - Train: {len(processed_data['train'])} cặp câu")
+    print(f"   - Validation: {len(processed_data['validation'])} cặp câu")
+    print(f"   - Test: {len(processed_data['test'])} cặp câu")
+    
+    # 5. Lưu thông tin tokenizer (để dùng sau)
+    tokenizer_info = {
+        'vocab_size': tokenizer.get_vocab_size(),
+        'sos_id': tokenizer.token_to_id("<sos>"),
+        'eos_id': tokenizer.token_to_id("<eos>"),
+        'pad_id': tokenizer.token_to_id("<pad>"),
+        'unk_id': tokenizer.token_to_id("<unk>"),
+        'tokenizer_path': str(tokenizer_path)
+    }
+    
+    info_path = PROCESSED_DATA_DIR / 'shared_vocab_info.json'
+    with open(info_path, 'w', encoding='utf-8') as f:
+        json.dump(tokenizer_info, f, indent=2)
+    
+    print(f"   ✓ Đã lưu tokenizer info: {info_path}")
+    
+    print("\n" + "="*70)
+    print("✓ HOÀN TẤT ENCODE DỮ LIỆU!")
+    print("="*70)
+    
+    print("\n📝 BƯỚC TIẾP THEO:")
+    print("   1. Cập nhật model để hỗ trợ shared vocabulary")
+    print("   2. Chạy training với shared vocab mode")
+
+# ============================================================================
+# RUN
+# ============================================================================
+
+if __name__ == "__main__":
+    main()
+

src/add_direction_tokens.py

@@ -0,0 +1,62 @@
+"""
+Thêm direction tokens (<en2vi> và <vi2en>) vào tokenizer hiện có
+"""
+
+from pathlib import Path
+from tokenizers import Tokenizer
+import json
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+PROCESSED_DATA_DIR = PROJECT_ROOT / 'data' / 'processed'
+TOKENIZER_PATH = PROCESSED_DATA_DIR / 'tokenizer_shared.json'
+
+def add_direction_tokens():
+    """Thêm direction tokens vào tokenizer"""
+    print("="*70)
+    print("THÊM DIRECTION TOKENS VÀO TOKENIZER")
+    print("="*70)
+    
+    # Load tokenizer hiện tại
+    if not TOKENIZER_PATH.exists():
+        raise FileNotFoundError(f"Không tìm thấy tokenizer tại {TOKENIZER_PATH}")
+    
+    print(f"\n📚 Đang load tokenizer từ: {TOKENIZER_PATH}")
+    tokenizer = Tokenizer.from_file(str(TOKENIZER_PATH))
+    
+    # Kiểm tra xem tokens đã có chưa
+    vocab = tokenizer.get_vocab()
+    has_en2vi = '<en2vi>' in vocab
+    has_vi2en = '<vi2en>' in vocab
+    
+    if has_en2vi and has_vi2en:
+        print("✓ Direction tokens đã có trong vocabulary!")
+        print(f"   <en2vi>: ID {vocab['<en2vi']}")
+        print(f"   <vi2en>: ID {vocab['<vi2en']}")
+        return tokenizer
+    
+    print("\n➕ Đang thêm direction tokens...")
+    
+    # Thêm tokens mới
+    # Lấy vocab size hiện tại
+    current_vocab_size = len(vocab)
+    
+    # Thêm tokens mới với ID tiếp theo
+    tokenizer.add_tokens(["<en2vi>", "<vi2en>"])
+    
+    # Save lại
+    print(f"\n💾 Đang lưu tokenizer với direction tokens...")
+    tokenizer.save(str(TOKENIZER_PATH))
+    
+    # Kiểm tra lại
+    vocab = tokenizer.get_vocab()
+    print(f"\n✓ Đã thêm direction tokens!")
+    print(f"   <en2vi>: ID {vocab.get('<en2vi>', 'NOT FOUND')}")
+    print(f"   <vi2en>: ID {vocab.get('<vi2en>', 'NOT FOUND')}")
+    print(f"   Vocab size: {len(vocab):,} (tăng từ {current_vocab_size:,})")
+    
+    return tokenizer
+
+if __name__ == '__main__':
+    add_direction_tokens()
+    print("\n✅ Hoàn thành!")
+

src/app.py

@@ -0,0 +1,413 @@
+"""
+GRADIO UI - VIETNAMESE TO ENGLISH TRANSLATION
+Modern web interface for Transformer translation model
+"""
+
+import glob
+import importlib
+import os
+import pickle
+import sys
+from pathlib import Path
+
+import gradio as gr
+import torch
+from shared_vocab_utils import (
+    create_shared_vocab_wrapper,
+    load_shared_vocab_info,
+)
+
+# ============================================================================
+# PATH CONFIGURATION
+# ============================================================================
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.append(str(PROJECT_ROOT))
+SRC_DIR = PROJECT_ROOT / 'src'
+if str(SRC_DIR) not in sys.path:
+    sys.path.append(str(SRC_DIR))
+
+DATA_PROCESSED_DIR = PROJECT_ROOT / 'data' / 'processed'
+CHECKPOINT_DIR = PROJECT_ROOT / 'checkpoints'
+
+# ============================================================================
+# DYNAMIC IMPORTS (giữ compatibility với pickle)
+# ============================================================================
+
+data_preprocessing = importlib.import_module('src.data_preprocessing')
+sys.modules.setdefault('data_preprocessing', data_preprocessing)
+Vocabulary = data_preprocessing.Vocabulary
+clean_text = data_preprocessing.clean_text
+
+complete_transformer = importlib.import_module('src.complete_transformer')
+create_model = complete_transformer.create_model
+
+inference_module = importlib.import_module('src.inference_evaluation')
+translate_sentence = inference_module.translate_sentence
+beam_search_decode = inference_module.beam_search_decode
+greedy_decode = inference_module.greedy_decode
+
+# ============================================================================
+# LOAD MODEL & VOCABULARIES
+# ============================================================================
+
+class TranslationModel:
+    """Wrapper class để quản lý model và vocabularies"""
+    
+    def __init__(self):
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.model = None
+        self.vi_vocab = None
+        self.en_vocab = None
+        self.model_info = {}
+        
+    def load_latest_model(self):
+        """Load model mới nhất từ checkpoints"""
+        print("🔄 Loading model và vocabularies...")
+
+        # Try shared vocabulary first (tokenizer_shared.json)
+        try:
+            vocab_info = load_shared_vocab_info()
+            self.vi_vocab, self.en_vocab = create_shared_vocab_wrapper()
+            self.pad_idx = vocab_info.get("pad_id", 0)
+            self.sos_idx = vocab_info.get("sos_id", 1)
+            self.eos_idx = vocab_info.get("eos_id", 2)
+            vocab_size = vocab_info.get("vocab_size")
+            print(f"✓ Loaded shared vocabulary (size={vocab_size})")
+        except Exception as shared_err:
+            print(f"⚠️ Shared vocab load failed, fallback to pickled vocabs: {shared_err}")
+            # Load legacy pickled vocabs
+            vi_vocab_path = DATA_PROCESSED_DIR / 'vi_vocab.pkl'
+            en_vocab_path = DATA_PROCESSED_DIR / 'en_vocab.pkl'
+            try:
+                with open(vi_vocab_path, 'rb') as f:
+                    self.vi_vocab = pickle.load(f)
+                with open(en_vocab_path, 'rb') as f:
+                    self.en_vocab = pickle.load(f)
+                self.pad_idx = getattr(self.vi_vocab, "PAD_IDX", 0)
+                vocab_size = len(self.vi_vocab)
+                print(f"✓ Loaded vocabularies (VI: {len(self.vi_vocab)}, EN: {len(self.en_vocab)})")
+            except Exception as e:
+                raise Exception(f"❌ Không thể load vocabularies (shared & pickled đều lỗi): {e}")
+        
+        # Tìm model mới nhất
+        checkpoint_dir = CHECKPOINT_DIR
+        best_model_path = checkpoint_dir / 'best_model.pt'
+        
+        # Ưu tiên best_model.pt
+        if os.path.exists(best_model_path):
+            checkpoint_path = best_model_path
+            print(f"✓ Found best model: {best_model_path}")
+        else:
+            # Tìm checkpoint mới nhất
+            checkpoints = glob.glob(str(checkpoint_dir / 'checkpoint_epoch_*.pt'))
+            if not checkpoints:
+                raise Exception(f"❌ Không tìm thấy checkpoint nào trong {checkpoint_dir}")
+            
+            # Sắp xếp theo thời gian modified
+            checkpoint_path = max(checkpoints, key=os.path.getmtime)
+            print(f"✓ Found latest checkpoint: {checkpoint_path}")
+        
+        # Load checkpoint để lấy thông tin
+        checkpoint = torch.load(checkpoint_path, map_location=self.device)
+        
+        # Tạo model
+        print("🔨 Creating model...")
+        # Detect model size từ checkpoint path (simple heuristic)
+        path_str = str(checkpoint_path).lower()
+        model_size = 'custom_25m'  # Default to trained config (d_model=384, 6 layers)
+        for size in ['custom_25m', 'medium', 'base', 'small', 'tiny', 'large']:
+            if size in path_str or ('25m' in path_str and size == 'custom_25m'):
+                model_size = size
+                break
+        
+        self.model, model_config = create_model(
+            src_vocab_size=len(self.vi_vocab),
+            tgt_vocab_size=len(self.en_vocab),
+            model_size=model_size,
+            pad_idx=getattr(self, "pad_idx", 0),
+            use_shared_vocab=True,
+            use_weight_tying=True,
+        )
+        
+        # Load weights
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        self.model = self.model.to(self.device)
+        self.model.eval()
+        
+        # Lưu thông tin model
+        self.model_info = {
+            'checkpoint': checkpoint_path,
+            'epoch': checkpoint.get('epoch', 'N/A'),
+            'val_loss': checkpoint.get('val_loss', 'N/A'),
+            'val_ppl': checkpoint.get('val_ppl', 'N/A'),
+            'model_size': model_size,
+            'device': str(self.device),
+            'parameters': sum(p.numel() for p in self.model.parameters()),
+            'vocab_size': len(self.vi_vocab),
+        }
+        
+        print(f"✓ Model loaded successfully!")
+        print(f"  - Epoch: {self.model_info['epoch']}")
+        print(f"  - Val Loss: {self.model_info['val_loss']}")
+        print(f"  - Device: {self.model_info['device']}")
+        
+        return self.model_info
+    
+    def translate(self, text, direction='vi2en', use_beam_search=True, beam_size=5):
+        """Dịch một câu theo hướng chỉ định"""
+        if not self.model or not self.vi_vocab or not self.en_vocab:
+            return "❌ Model chưa được load. Vui lòng reload model."
+        
+        try:
+            # Xác định ngôn ngữ source
+            src_lang = 'en' if direction == 'en2vi' else 'vi'
+            
+            # Clean text
+            text = clean_text(text, src_lang)
+            
+            if not text.strip():
+                lang_name = "tiếng Anh" if direction == 'en2vi' else "tiếng Việt"
+                return f"⚠️ Vui lòng nhập văn bản {lang_name}."
+            
+            # Với shared vocabulary, cả vi_vocab và en_vocab đều là cùng một tokenizer
+            # Model có thể dịch cả hai chiều vì dùng shared vocabulary
+            # Translate với shared vocab (vi_vocab và en_vocab giống nhau)
+            translation = translate_sentence(
+                self.model, text, self.vi_vocab, self.en_vocab,
+                self.device, use_beam_search, beam_size, 
+                src_lang=src_lang,
+                repetition_penalty=1.3,
+                no_repeat_ngram_size=3
+            )
+            
+            return translation
+        
+        except Exception as e:
+            import traceback
+            error_detail = traceback.format_exc()
+            return f"❌ Lỗi khi dịch: {str(e)}\n\nChi tiết:\n{error_detail}"
+
+# Khởi tạo model
+translator = TranslationModel()
+
+# ============================================================================
+# GRADIO INTERFACE
+# ============================================================================
+
+def translate_text(text, direction, use_beam_search, beam_size):
+    """Wrapper function cho Gradio"""
+    return translator.translate(text, direction, use_beam_search, beam_size)
+
+def reload_model():
+    """Reload model mới nhất"""
+    try:
+        info = translator.load_latest_model()
+        return f"""
+✅ **Model đã được load thành công!**
+
+📊 **Thông tin model:**
+- Checkpoint: `{info['checkpoint']}`
+- Epoch: {info['epoch']}
+- Validation Loss: {f"{info['val_loss']:.4f}" if isinstance(info['val_loss'], float) else info['val_loss']}
+- Model Size: {info['model_size'].upper()}
+- Device: {info['device']}
+- Parameters: {info['parameters']:,}
+- Vocab size: {info.get('vocab_size', 'N/A')}
+"""
+    except Exception as e:
+        return f"❌ Lỗi khi load model: {str(e)}"
+
+def get_model_info():
+    """Lấy thông tin model hiện tại"""
+    if not translator.model:
+        return "⚠️ Model chưa được load. Click 'Reload Model' để load model."
+    
+    info = translator.model_info
+    return f"""
+📊 **Model hiện tại:**
+- Checkpoint: `{info.get('checkpoint', 'N/A')}`
+- Epoch: {info.get('epoch', 'N/A')}
+- Validation Loss: {f"{info['val_loss']:.4f}" if isinstance(info.get('val_loss'), float) else info.get('val_loss', 'N/A')}
+- Model Size: {str(info.get('model_size', 'N/A')).upper()}
+- Device: {info.get('device', 'N/A')}
+- Parameters: {info.get('parameters', 0):,}
+- Vocab size: {info.get('vocab_size', 'N/A')}
+"""
+
+# ============================================================================
+# BUILD GRADIO APP
+# ============================================================================
+
+def create_app():
+    """Tạo Gradio app"""
+    
+    with gr.Blocks(
+        title="Bidirectional Translation: English ↔ Vietnamese"
+    ) as app:
+        
+        # Header
+        gr.Markdown("""
+        # 🌐 Bidirectional Translation: English ↔ Vietnamese
+        ### Powered by Transformer Neural Network
+        
+        Dịch văn bản hai chiều giữa **Tiếng Anh** và **Tiếng Việt** sử dụng mô hình Transformer được train từ đầu.
+        """)
+        
+        # Main content
+        with gr.Row():
+            with gr.Column(scale=1):
+                # Direction selector
+                gr.Markdown("### 🔄 Chọn hướng dịch")
+                direction = gr.Radio(
+                    choices=[
+                        ("Vietnamese → English", "vi2en"),
+                        ("English → Vietnamese", "en2vi")
+                    ],
+                    value="vi2en",
+                    label="Hướng dịch",
+                    info="Chọn ngôn ngữ đầu vào và đầu ra"
+                )
+                
+                # Input
+                gr.Markdown("### 📝 Nhập văn bản")
+                input_text = gr.Textbox(
+                    label="Văn bản đầu vào",
+                    placeholder="Nhập câu cần dịch...\nVí dụ: Xin chào, tôi là sinh viên.\nHoặc: Hello, I am a student.",
+                    lines=5,
+                    max_lines=10
+                )
+                
+                # Options
+                with gr.Row():
+                    use_beam_search = gr.Checkbox(
+                        label="🔍 Beam Search",
+                        value=True,
+                        info="Kết quả tốt hơn nhưng chậm hơn"
+                    )
+                    beam_size = gr.Slider(
+                        minimum=1,
+                        maximum=10,
+                        value=5,
+                        step=1,
+                        label="Beam Size",
+                        info="Số candidates (cao hơn = tốt hơn nhưng chậm hơn)"
+                    )
+                
+                # Buttons
+                with gr.Row():
+                    translate_btn = gr.Button("🚀 Dịch", variant="primary", size="lg")
+                    clear_btn = gr.ClearButton([input_text], value="🗑️ Xóa")
+            
+            with gr.Column(scale=1):
+                # Output
+                gr.Markdown("### 🎯 Kết quả dịch")
+                output_text = gr.Textbox(
+                    label="Văn bản đã dịch",
+                    placeholder="Kết quả dịch sẽ hiển thị ở đây...",
+                    lines=5,
+                    max_lines=10,
+                    interactive=False
+                )
+        
+        # Examples
+        gr.Markdown("### 💡 Ví dụ")
+        gr.Examples(
+            examples=[
+                ["vi2en", "Xin chào, tôi tên là Nam.", True, 5],
+                ["vi2en", "Hôm nay thời tiết đẹp.", True, 5],
+                ["vi2en", "Tôi đang học tiếng Anh.", True, 5],
+                ["vi2en", "Bạn khỏe không?", True, 5],
+                ["en2vi", "Hello, my name is Nam.", True, 5],
+                ["en2vi", "The weather is beautiful today.", True, 5],
+                ["en2vi", "I am learning English.", True, 5],
+                ["en2vi", "How are you?", True, 5],
+            ],
+            inputs=[direction, input_text, use_beam_search, beam_size],
+            outputs=output_text,
+            fn=translate_text,
+            cache_examples=False,
+        )
+        
+        # Model info section
+        with gr.Accordion("⚙️ Thông tin Model & Cài đặt", open=False):
+            model_info_display = gr.Markdown(get_model_info())
+            
+            with gr.Row():
+                reload_btn = gr.Button("🔄 Reload Model", variant="secondary")
+                info_btn = gr.Button("ℹ️ Xem thông tin", variant="secondary")
+            
+            reload_output = gr.Markdown()
+        
+        # Connect buttons
+        translate_btn.click(
+            fn=translate_text,
+            inputs=[input_text, direction, use_beam_search, beam_size],
+            outputs=output_text
+        )
+        
+        reload_btn.click(
+            fn=reload_model,
+            outputs=reload_output
+        )
+        
+        info_btn.click(
+            fn=get_model_info,
+            outputs=model_info_display
+        )
+        
+        # Footer
+        gr.Markdown("""
+        ---
+        ### 📚 Hướng dẫn sử dụng
+        1. **Chọn hướng dịch**: Vietnamese → English hoặc English → Vietnamese
+        2. **Nhập văn bản** vào ô bên trái
+        3. **Chọn phương pháp dịch**:
+           - ✅ Beam Search: Chất lượng cao hơn (khuyên dùng)
+           - ❌ Greedy Search: Nhanh hơn nhưng kém chính xác hơn
+        4. **Click "Dịch"** để xem kết quả
+        5. **Reload Model** nếu bạn vừa train xong model mới
+        
+        ### 🔧 Tips
+        - Model hỗ trợ dịch hai chiều: Việt ↔ Anh
+        - Beam size càng cao thì kết quả càng tốt nhưng chậm hơn (khuyên dùng 5)
+        - Câu ngắn dịch nhanh hơn câu dài
+        - Model hoạt động tốt nhất với câu có độ dài 5-20 từ
+        - Model sử dụng shared vocabulary và checkpoint hiện có, không cần train lại
+        """)
+    
+    return app
+
+# ============================================================================
+# LAUNCH APP
+# ============================================================================
+
+if __name__ == "__main__":
+    print("="*70)
+    print("KHỞI ĐỘNG TRANSLATION WEB APP")
+    print("="*70)
+    
+    # Load model khi khởi động
+    try:
+        translator.load_latest_model()
+    except Exception as e:
+        print(f"⚠️ Warning: {e}")
+        print("💡 Bạn có thể reload model sau trong giao diện web.")
+    
+    # Create and launch app
+    app = create_app()
+    
+    print("\n" + "="*70)
+    print("🚀 LAUNCHING WEB APP...")
+    print("="*70)
+    
+    import os
+    server_port = int(os.environ.get("GRADIO_SERVER_PORT", "7860"))
+    
+    app.launch(
+        server_name="0.0.0.0",  # Cho phép truy cập từ máy khác
+        server_port=server_port,
+        share=False,  # Set True nếu muốn share link public
+        inbrowser=True,  # Tự động mở browser
+    )

src/bidirectional_translate.py

@@ -0,0 +1,546 @@
+"""
+BIDIRECTIONAL TRANSLATION SCRIPT
+Dịch hai chiều: English ↔ Vietnamese
+Sử dụng checkpoint và vocab hiện có, không cần train lại
+"""
+
+import torch
+import argparse
+from pathlib import Path
+from tqdm import tqdm
+import sys
+
+# Add parent directory to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from src.complete_transformer import TransformerShared
+from src.shared_vocab_utils import load_shared_tokenizer, load_shared_vocab_info, clean_decoded_output
+from src.inference_evaluation import beam_search_decode, greedy_decode
+from src.data_preprocessing import clean_text
+
+# ============================================================================
+# CONFIGURATION
+# ============================================================================
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DATA_DIR = PROJECT_ROOT / 'data'
+RESULTS_DIR = PROJECT_ROOT / 'results'
+CHECKPOINT_DIR = PROJECT_ROOT / 'checkpoints'
+
+# ============================================================================
+# LOAD MODEL
+# ============================================================================
+
+def load_model(checkpoint_path, model_size='custom_25m', device='cuda'):
+    """
+    Load model từ checkpoint
+    
+    Args:
+        checkpoint_path: Đường dẫn đến checkpoint file
+        model_size: Kích thước model ('custom_25m', 'base', etc.)
+        device: Device để load model
+        
+    Returns:
+        model: Loaded model
+        vocab_info: Vocabulary info dict
+    """
+    print(f"\n{'='*70}")
+    print(f"LOADING MODEL")
+    print(f"{'='*70}")
+    
+    # Load vocab info
+    vocab_info = load_shared_vocab_info()
+    vocab_size = vocab_info['vocab_size']
+    
+    print(f"Vocab size: {vocab_size}")
+    
+    # Model configs theo model_size
+    model_configs = {
+        'custom_25m': {
+            'd_model': 384,
+            'n_layers': 6,
+            'n_heads': 8,
+            'd_ff': 1536,
+            'dropout': 0.1
+        },
+        'base': {
+            'd_model': 512,
+            'n_layers': 6,
+            'n_heads': 8,
+            'd_ff': 2048,
+            'dropout': 0.1
+        }
+    }
+    
+    config = model_configs.get(model_size, model_configs['custom_25m'])
+    
+    # Tạo model
+    model = TransformerShared(
+        vocab_size=vocab_size,
+        d_model=config['d_model'],
+        n_layers=config['n_layers'],
+        n_heads=config['n_heads'],
+        d_ff=config['d_ff'],
+        dropout=config['dropout'],
+        pad_idx=vocab_info['pad_id']
+    )
+    
+    # Load checkpoint
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model = model.to(device)
+    model.eval()
+    
+    print(f"✓ Model loaded from {checkpoint_path}")
+    print(f"✓ Model size: {model_size}")
+    print(f"✓ Device: {device}")
+    print(f"{'='*70}\n")
+    
+    return model, vocab_info
+
+# ============================================================================
+# VOCABULARY WRAPPER
+# ============================================================================
+
+class VocabWrapper:
+    """Wrapper cho tokenizer để tương thích với inference functions"""
+    def __init__(self, tokenizer, vocab_info):
+        self.tokenizer = tokenizer
+        self.SOS_IDX = vocab_info['sos_id']
+        self.EOS_IDX = vocab_info['eos_id']
+        self.PAD_IDX = vocab_info['pad_id']
+        self.UNK_IDX = vocab_info['unk_id']
+    
+    def encode(self, text):
+        """Encode text thành token IDs"""
+        encoding = self.tokenizer.encode(text)
+        return encoding.ids
+    
+    def decode(self, token_ids):
+        """Decode token IDs thành text (nhận list of token IDs)"""
+        # Filter out special tokens
+        filtered_ids = [tid for tid in token_ids 
+                       if tid not in [self.SOS_IDX, self.EOS_IDX, self.PAD_IDX]]
+        if not filtered_ids:
+            return ""
+        # Decode using tokenizer
+        decoded = self.tokenizer.decode(filtered_ids)
+        # Post-processing: Clean output
+        decoded = clean_decoded_output(decoded)
+        return decoded
+
+# ============================================================================
+# POST-PROCESSING: CLEAN TRANSLATION OUTPUT
+# ============================================================================
+
+def clean_translation_output(text):
+    """
+    Làm sạch output translation để loại bỏ rác ở cuối câu
+    
+    Args:
+        text: Text cần làm sạch
+    
+    Returns:
+        Text đã được làm sạch
+    """
+    if not text:
+        return text
+    
+    import re
+    
+    # Loại bỏ các pattern rác ở cuối câu
+    # Pattern: – số., – ", – [ [ [, etc.
+    text = re.sub(r'\s*–\s*\d+\.?\s*$', '', text)  # – 7., – 1.
+    text = re.sub(r'\s*–\s*["\']\s*$', '', text)  # – ", – '
+    text = re.sub(r'\s*–\s*\[\s*\[\s*\[\s*$', '', text)  # – [ [ [
+    text = re.sub(r'\s*–\s*$', '', text)  # – ở cuối
+    text = re.sub(r'\s*–\s*[^\w\s.,!?;:()\[\]{}"\'-]+$', '', text)  # – ký tự đặc biệt
+    
+    # Loại bỏ số trang và footnote ở cuối
+    text = re.sub(r'\s+\d+\.\s*$', '', text)  # " 7. "
+    text = re.sub(r'\s+\[\d+\]\s*$', '', text)  # " [1] "
+    
+    # Loại bỏ các ký tự đặc biệt thừa ở cuối (nhưng giữ dấu câu hợp lệ)
+    text = re.sub(r'[^\w\s.,!?;:()\[\]{}"\'-]+$', '', text)
+    
+    # Loại bỏ khoảng trắng thừa
+    text = ' '.join(text.split())
+    
+    # Loại bỏ lặp từ ở cuối (nếu có pattern như "word word word")
+    words = text.split()
+    if len(words) >= 3:
+        # Kiểm tra 3 từ cuối có lặp không
+        last_3 = words[-3:]
+        if len(set(last_3)) == 1:  # Tất cả giống nhau
+            text = ' '.join(words[:-2])  # Loại bỏ 2 từ cuối
+    
+    return text.strip()
+
+# ============================================================================
+# TRANSLATION FUNCTIONS
+# ============================================================================
+
+def translate_sentence_bidirectional(
+    model,
+    sentence,
+    direction,
+    tokenizer,
+    vocab_info,
+    device='cuda',
+    use_beam_search=True,
+    beam_size=5,
+    max_len=100,
+    repetition_penalty=1.3,
+    no_repeat_ngram_size=3
+):
+    """
+    Dịch một câu theo hướng chỉ định
+    
+    Args:
+        model: Transformer model
+        sentence: Source sentence (string)
+        direction: 'en2vi' hoặc 'vi2en'
+        tokenizer: Tokenizer
+        vocab_info: Vocabulary info dict
+        device: Device
+        use_beam_search: Dùng beam search hay greedy
+        beam_size: Beam size
+        max_len: Max decode length
+        repetition_penalty: Penalty cho tokens lặp lại
+        no_repeat_ngram_size: Kích thước n-gram để tránh lặp lại
+        
+    Returns:
+        translation: Translated sentence
+    """
+    model.eval()
+    
+    # Xác định ngôn ngữ source để clean text
+    src_lang = 'en' if direction == 'en2vi' else 'vi'
+    
+    # Tiền xử lý câu
+    sentence = clean_text(sentence, src_lang)
+    
+    if not sentence.strip():
+        return ""
+    
+    # Tạo vocab wrapper
+    vocab = VocabWrapper(tokenizer, vocab_info)
+    
+    # Lấy direction token ID
+    EN2VI_ID = tokenizer.token_to_id("<en2vi>")
+    VI2EN_ID = tokenizer.token_to_id("<vi2en>")
+    SOS_ID = vocab_info['sos_id']
+    
+    if EN2VI_ID is None or VI2EN_ID is None:
+        raise ValueError("Không tìm thấy direction tokens trong tokenizer!")
+    
+    direction_id = EN2VI_ID if direction == 'en2vi' else VI2EN_ID
+    
+    # Encode
+    try:
+        tokens = vocab.encode(sentence)
+        if not tokens:
+            return ""
+        # Thêm SOS và direction token vào đầu
+        tokens = [SOS_ID, direction_id] + tokens
+    except Exception as e:
+        print(f"⚠️ Lỗi khi encode: {e}")
+        return ""
+    
+    src = torch.LongTensor([tokens]).to(device)
+    
+    # Decode với shared vocabulary (cùng vocab cho cả source và target)
+    if use_beam_search:
+        _, translation = beam_search_decode(
+            model, src, vocab, vocab, 
+            device, beam_size, max_len, 
+            alpha=0.6, 
+            repetition_penalty=repetition_penalty,
+            no_repeat_ngram_size=no_repeat_ngram_size
+        )
+    else:
+        _, translation = greedy_decode(
+            model, src, vocab, vocab, 
+            device, max_len,
+            repetition_penalty=repetition_penalty,
+            no_repeat_ngram_size=no_repeat_ngram_size
+        )
+    
+    return translation
+
+def translate_file_bidirectional(
+    model,
+    tokenizer,
+    vocab_info,
+    source_file,
+    output_file,
+    direction='vi2en',
+    device='cuda',
+    use_beam_search=True,
+    beam_size=5,
+    max_len=100,
+    repetition_penalty=1.3,
+    no_repeat_ngram_size=3
+):
+    """
+    Dịch file source và lưu kết quả
+    
+    Args:
+        model: Transformer model
+        tokenizer: Tokenizer
+        vocab_info: Vocabulary info dict
+        source_file: File chứa source sentences (mỗi dòng 1 câu)
+        output_file: File output để lưu translations
+        direction: 'en2vi' hoặc 'vi2en'
+        device: Device
+        use_beam_search: Dùng beam search hay greedy
+        beam_size: Beam size
+        max_len: Max decode length
+    """
+    print(f"\n{'='*70}")
+    print(f"INFERENCE - {direction.upper()}")
+    print(f"{'='*70}")
+    print(f"Source file: {source_file}")
+    print(f"Output file: {output_file}")
+    print(f"Direction: {'English → Vietnamese' if direction == 'en2vi' else 'Vietnamese → English'}")
+    print(f"Method: {'Beam Search' if use_beam_search else 'Greedy Search'}")
+    if use_beam_search:
+        print(f"Beam size: {beam_size}")
+    print(f"{'='*70}\n")
+    
+    # Load source sentences
+    with open(source_file, 'r', encoding='utf-8') as f:
+        source_sentences = [line.strip() for line in f.readlines()]
+    
+    print(f"Loaded {len(source_sentences)} source sentences")
+    
+    # Translate
+    translations = []
+    model.eval()
+    
+    with torch.no_grad():
+        for src_text in tqdm(source_sentences, desc='Translating'):
+            translation = translate_sentence_bidirectional(
+                model=model,
+                sentence=src_text,
+                direction=direction,
+                tokenizer=tokenizer,
+                vocab_info=vocab_info,
+                device=device,
+                use_beam_search=use_beam_search,
+                beam_size=beam_size,
+                max_len=max_len,
+                repetition_penalty=repetition_penalty,
+                no_repeat_ngram_size=no_repeat_ngram_size
+            )
+            # Post-process: Loại bỏ rác ở cuối câu
+            translation = clean_translation_output(translation)
+            translations.append(translation)
+    
+    # Save translations
+    RESULTS_DIR.mkdir(exist_ok=True, parents=True)
+    with open(output_file, 'w', encoding='utf-8') as f:
+        for trans in translations:
+            f.write(trans + '\n')
+    
+    print(f"\n✓ Saved {len(translations)} translations to {output_file}")
+    print(f"{'='*70}\n")
+
+def interactive_translation_bidirectional(
+    model,
+    tokenizer,
+    vocab_info,
+    device='cuda',
+    use_beam_search=True,
+    beam_size=5,
+    repetition_penalty=1.3,
+    no_repeat_ngram_size=3
+):
+    """
+    Chế độ dịch tương tác hai chiều
+    
+    Args:
+        model: Transformer model
+        tokenizer: Tokenizer
+        vocab_info: Vocabulary info dict
+        device: Device
+        use_beam_search: Sử dụng beam search
+        beam_size: Beam size
+    """
+    print("\n" + "="*70)
+    print("CHẾ ĐỘ DỊCH TƯƠNG TÁC HAI CHIỀU")
+    print("="*70)
+    print("Hướng dẫn:")
+    print("  - Nhập 'en:' hoặc 'vi:' để chỉ định ngôn ngữ đầu vào")
+    print("  - Ví dụ: 'en: Hello, how are you?' → dịch sang tiếng Việt")
+    print("  - Ví dụ: 'vi: Xin chào, bạn khỏe không?' → dịch sang tiếng Anh")
+    print("  - Gõ 'quit' hoặc 'exit' để thoát")
+    print("="*70 + "\n")
+    
+    while True:
+        user_input = input("Nhập câu (hoặc 'en:...' / 'vi:...'): ").strip()
+        
+        if user_input.lower() in ['quit', 'exit', '']:
+            print("Tạm biệt!")
+            break
+        
+        # Xác định hướng dịch
+        if user_input.startswith('en:'):
+            direction = 'en2vi'
+            sentence = user_input[3:].strip()
+            target_lang = "Tiếng Việt"
+        elif user_input.startswith('vi:'):
+            direction = 'vi2en'
+            sentence = user_input[3:].strip()
+            target_lang = "Tiếng Anh"
+        else:
+            # Mặc định: tự động phát hiện (đơn giản)
+            # Nếu có nhiều ký tự Latin không dấu → tiếng Anh
+            # Ngược lại → tiếng Việt
+            has_vietnamese_chars = any(ord(c) >= 0x0100 for c in user_input)
+            if has_vietnamese_chars:
+                direction = 'vi2en'
+                target_lang = "Tiếng Anh"
+            else:
+                direction = 'en2vi'
+                target_lang = "Tiếng Việt"
+            sentence = user_input
+        
+        if not sentence:
+            print("⚠️ Vui lòng nhập câu cần dịch.\n")
+            continue
+        
+        try:
+            translation = translate_sentence_bidirectional(
+                model=model,
+                sentence=sentence,
+                direction=direction,
+                tokenizer=tokenizer,
+                vocab_info=vocab_info,
+                device=device,
+                use_beam_search=use_beam_search,
+                beam_size=beam_size,
+                repetition_penalty=repetition_penalty,
+                no_repeat_ngram_size=no_repeat_ngram_size
+            )
+            print(f"{target_lang}: {translation}\n")
+        except Exception as e:
+            print(f"❌ Lỗi khi dịch: {e}\n")
+
+# ============================================================================
+# MAIN
+# ============================================================================
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Dịch hai chiều: English ↔ Vietnamese',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Ví dụ sử dụng:
+  # Dịch file tiếng Việt sang tiếng Anh
+  python src/bidirectional_translate.py --checkpoint checkpoints/best_model.pt \\
+      --source data/raw/test.vi --output results/test_en.txt --direction vi2en
+  
+  # Dịch file tiếng Anh sang tiếng Việt
+  python src/bidirectional_translate.py --checkpoint checkpoints/best_model.pt \\
+      --source data/raw/test.en --output results/test_vi.txt --direction en2vi
+  
+  # Chế độ tương tác
+  python src/bidirectional_translate.py --checkpoint checkpoints/best_model.pt --interactive
+        """
+    )
+    parser.add_argument('--checkpoint', type=str, required=True,
+                       help='Đường dẫn đến checkpoint file')
+    parser.add_argument('--source', type=str, default=None,
+                       help='File source (mỗi dòng 1 câu)')
+    parser.add_argument('--output', type=str, default=None,
+                       help='File output để lưu translations')
+    parser.add_argument('--direction', type=str, default='vi2en',
+                       choices=['vi2en', 'en2vi'],
+                       help='Hướng dịch: vi2en (Việt→Anh) hoặc en2vi (Anh→Việt)')
+    parser.add_argument('--model_size', type=str, default='custom_25m',
+                       choices=['custom_25m', 'base'],
+                       help='Kích thước model')
+    parser.add_argument('--beam_size', type=int, default=5,
+                       help='Beam size cho beam search')
+    parser.add_argument('--greedy', action='store_true',
+                       help='Dùng greedy search thay vì beam search')
+    parser.add_argument('--max_len', type=int, default=100,
+                       help='Maximum decode length')
+    parser.add_argument('--repetition_penalty', type=float, default=1.3,
+                       help='Repetition penalty factor (>1.0 để giảm repetition)')
+    parser.add_argument('--no_repeat_ngram_size', type=int, default=3,
+                       help='Kích thước n-gram để tránh lặp lại (0 = tắt)')
+    parser.add_argument('--device', type=str, default='cuda',
+                       help='Device (cuda/cpu)')
+    parser.add_argument('--interactive', action='store_true',
+                       help='Chế độ dịch tương tác')
+    
+    args = parser.parse_args()
+    
+    # Check checkpoint
+    checkpoint_path = Path(args.checkpoint)
+    if not checkpoint_path.exists():
+        print(f"❌ Không tìm thấy checkpoint: {checkpoint_path}")
+        return
+    
+    # Device
+    device = args.device
+    if device == 'cuda' and not torch.cuda.is_available():
+        print("⚠️  CUDA không khả dụng, dùng CPU")
+        device = 'cpu'
+    
+    # Load model
+    model, vocab_info = load_model(
+        checkpoint_path, 
+        args.model_size, 
+        device
+    )
+    
+    # Load tokenizer
+    tokenizer = load_shared_tokenizer()
+    
+    # Interactive mode
+    if args.interactive:
+        interactive_translation_bidirectional(
+            model=model,
+            tokenizer=tokenizer,
+            vocab_info=vocab_info,
+            device=device,
+            use_beam_search=not args.greedy,
+            beam_size=args.beam_size,
+            repetition_penalty=args.repetition_penalty,
+            no_repeat_ngram_size=args.no_repeat_ngram_size
+        )
+        return
+    
+    # File translation mode
+    if not args.source or not args.output:
+        print("❌ Cần cung cấp --source và --output (hoặc dùng --interactive)")
+        parser.print_help()
+        return
+    
+    source_path = Path(args.source)
+    if not source_path.exists():
+        print(f"❌ Không tìm thấy source file: {source_path}")
+        return
+    
+    # Translate
+    translate_file_bidirectional(
+        model=model,
+        tokenizer=tokenizer,
+        vocab_info=vocab_info,
+        source_file=source_path,
+        output_file=Path(args.output),
+        direction=args.direction,
+        device=device,
+        use_beam_search=not args.greedy,
+        beam_size=args.beam_size,
+        max_len=args.max_len,
+        repetition_penalty=args.repetition_penalty,
+        no_repeat_ngram_size=args.no_repeat_ngram_size
+    )
+
+if __name__ == '__main__':
+    main()
+

src/checkpoint_manager.py

@@ -0,0 +1,251 @@
+"""
+CHECKPOINT MANAGER
+Tiện ích để quản lý và chuyển đổi giữa các checkpoint
+"""
+
+import torch
+import argparse
+import json
+from pathlib import Path
+from datetime import datetime
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+CHECKPOINT_DIR = PROJECT_ROOT / 'checkpoints'
+
+def list_checkpoints():
+    """Liệt kê tất cả checkpoint có sẵn"""
+    print("="*70)
+    print("DANH SÁCH CHECKPOINT")
+    print("="*70)
+    
+    checkpoints = list(CHECKPOINT_DIR.glob('*.pt'))
+    
+    if not checkpoints:
+        print("❌ Không tìm thấy checkpoint nào!")
+        return
+    
+    # Phân loại checkpoint
+    best_models = []
+    epoch_checkpoints = []
+    finetune_checkpoints = []
+    other_checkpoints = []
+    
+    for ckpt in sorted(checkpoints):
+        name = ckpt.name
+        if name == 'best_model.pt':
+            best_models.append(('best_model.pt', 'Original best model'))
+        elif name == 'best_model_finetuned.pt':
+            best_models.append(('best_model_finetuned.pt', 'Finetuned best model'))
+        elif name.startswith('checkpoint_epoch_'):
+            epoch_checkpoints.append(ckpt)
+        elif name.startswith('finetune_checkpoint_epoch_'):
+            finetune_checkpoints.append(ckpt)
+        else:
+            other_checkpoints.append(ckpt)
+    
+    # In best models
+    if best_models:
+        print("\n📌 BEST MODELS:")
+        for name, desc in best_models:
+            ckpt_path = CHECKPOINT_DIR / name
+            if ckpt_path.exists():
+                try:
+                    ckpt = torch.load(ckpt_path, map_location='cpu')
+                    epoch = ckpt.get('epoch', 'N/A')
+                    val_loss = ckpt.get('val_loss', 'N/A')
+                    size_mb = ckpt_path.stat().st_size / (1024 * 1024)
+                    print(f"  ✓ {name}")
+                    print(f"    - {desc}")
+                    print(f"    - Epoch: {epoch}")
+                    print(f"    - Val Loss: {val_loss:.4f}" if isinstance(val_loss, float) else f"    - Val Loss: {val_loss}")
+                    print(f"    - Size: {size_mb:.2f} MB")
+                    
+                    # Kiểm tra finetune config
+                    if 'finetune_config' in ckpt:
+                        finetune_cfg = ckpt['finetune_config']
+                        print(f"    - Finetuned from: {Path(finetune_cfg.get('checkpoint_path', 'N/A')).name}")
+                        print(f"    - LR Factor: {finetune_cfg.get('finetune_lr_factor', 'N/A')}")
+                    print()
+                except Exception as e:
+                    print(f"  ⚠️  {name} (lỗi khi đọc: {e})")
+    
+    # In epoch checkpoints
+    if epoch_checkpoints:
+        print(f"\n📦 EPOCH CHECKPOINTS ({len(epoch_checkpoints)} files):")
+        for ckpt in epoch_checkpoints[:5]:  # Chỉ hiển thị 5 đầu
+            try:
+                ckpt_data = torch.load(ckpt, map_location='cpu')
+                epoch = ckpt_data.get('epoch', 'N/A')
+                val_loss = ckpt_data.get('val_loss', 'N/A')
+                size_mb = ckpt.stat().st_size / (1024 * 1024)
+                print(f"  - {ckpt.name}: Epoch {epoch}, Val Loss: {val_loss:.4f}" if isinstance(val_loss, float) else f"  - {ckpt.name}: Epoch {epoch}, Val Loss: {val_loss}")
+                print(f"    Size: {size_mb:.2f} MB")
+            except:
+                print(f"  - {ckpt.name} (lỗi khi đọc)")
+        
+        if len(epoch_checkpoints) > 5:
+            print(f"  ... và {len(epoch_checkpoints) - 5} checkpoint khác")
+    
+    # In finetune checkpoints
+    if finetune_checkpoints:
+        print(f"\n🔧 FINETUNE CHECKPOINTS ({len(finetune_checkpoints)} files):")
+        for ckpt in finetune_checkpoints[:5]:  # Chỉ hiển thị 5 đầu
+            try:
+                ckpt_data = torch.load(ckpt, map_location='cpu')
+                epoch = ckpt_data.get('epoch', 'N/A')
+                val_loss = ckpt_data.get('val_loss', 'N/A')
+                size_mb = ckpt.stat().st_size / (1024 * 1024)
+                print(f"  - {ckpt.name}: Epoch {epoch}, Val Loss: {val_loss:.4f}" if isinstance(val_loss, float) else f"  - {ckpt.name}: Epoch {epoch}, Val Loss: {val_loss}")
+                print(f"    Size: {size_mb:.2f} MB")
+            except:
+                print(f"  - {ckpt.name} (lỗi khi đọc)")
+        
+        if len(finetune_checkpoints) > 5:
+            print(f"  ... và {len(finetune_checkpoints) - 5} checkpoint khác")
+    
+    # In other checkpoints
+    if other_checkpoints:
+        print(f"\n📁 OTHER CHECKPOINTS ({len(other_checkpoints)} files):")
+        for ckpt in other_checkpoints:
+            size_mb = ckpt.stat().st_size / (1024 * 1024)
+            print(f"  - {ckpt.name} ({size_mb:.2f} MB)")
+    
+    print("\n" + "="*70)
+
+def switch_checkpoint(source_name, target_name='best_model.pt', backup=True):
+    """
+    Chuyển đổi checkpoint (copy checkpoint thành best_model.pt)
+    
+    Args:
+        source_name: Tên checkpoint nguồn
+        target_name: Tên checkpoint đích (mặc định: best_model.pt)
+        backup: Có backup checkpoint cũ không
+    """
+    source_path = CHECKPOINT_DIR / source_name
+    target_path = CHECKPOINT_DIR / target_name
+    
+    if not source_path.exists():
+        print(f"❌ Không tìm thấy checkpoint: {source_path}")
+        return False
+    
+    # Backup checkpoint cũ nếu có
+    if target_path.exists() and backup:
+        backup_name = f"{target_name}.backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+        backup_path = CHECKPOINT_DIR / backup_name
+        print(f"→ Đang backup checkpoint cũ: {backup_name}")
+        import shutil
+        shutil.copy2(target_path, backup_path)
+        print(f"  ✓ Đã backup: {backup_path}")
+    
+    # Copy checkpoint mới
+    print(f"→ Đang copy {source_name} → {target_name}")
+    import shutil
+    shutil.copy2(source_path, target_path)
+    print(f"  ✓ Đã chuyển đổi checkpoint!")
+    
+    # Hiển thị thông tin checkpoint mới
+    try:
+        ckpt = torch.load(target_path, map_location='cpu')
+        epoch = ckpt.get('epoch', 'N/A')
+        val_loss = ckpt.get('val_loss', 'N/A')
+        print(f"\n  Thông tin checkpoint mới:")
+        print(f"    - Epoch: {epoch}")
+        print(f"    - Val Loss: {val_loss:.4f}" if isinstance(val_loss, float) else f"    - Val Loss: {val_loss}")
+        
+        if 'finetune_config' in ckpt:
+            finetune_cfg = ckpt['finetune_config']
+            print(f"    - Finetuned from: {Path(finetune_cfg.get('checkpoint_path', 'N/A')).name}")
+    except Exception as e:
+        print(f"  ⚠️  Lỗi khi đọc checkpoint: {e}")
+    
+    return True
+
+def compare_checkpoints(checkpoint1_name, checkpoint2_name):
+    """So sánh 2 checkpoint"""
+    ckpt1_path = CHECKPOINT_DIR / checkpoint1_name
+    ckpt2_path = CHECKPOINT_DIR / checkpoint2_name
+    
+    if not ckpt1_path.exists():
+        print(f"❌ Không tìm thấy checkpoint: {checkpoint1_name}")
+        return
+    
+    if not ckpt2_path.exists():
+        print(f"❌ Không tìm thấy checkpoint: {checkpoint2_name}")
+        return
+    
+    print("="*70)
+    print("SO SÁNH CHECKPOINT")
+    print("="*70)
+    
+    try:
+        ckpt1 = torch.load(ckpt1_path, map_location='cpu')
+        ckpt2 = torch.load(ckpt2_path, map_location='cpu')
+        
+        print(f"\n📌 {checkpoint1_name}:")
+        print(f"  - Epoch: {ckpt1.get('epoch', 'N/A')}")
+        print(f"  - Val Loss: {ckpt1.get('val_loss', 'N/A'):.4f}" if isinstance(ckpt1.get('val_loss'), float) else f"  - Val Loss: {ckpt1.get('val_loss', 'N/A')}")
+        print(f"  - Size: {ckpt1_path.stat().st_size / (1024 * 1024):.2f} MB")
+        
+        print(f"\n📌 {checkpoint2_name}:")
+        print(f"  - Epoch: {ckpt2.get('epoch', 'N/A')}")
+        print(f"  - Val Loss: {ckpt2.get('val_loss', 'N/A'):.4f}" if isinstance(ckpt2.get('val_loss'), float) else f"  - Val Loss: {ckpt2.get('val_loss', 'N/A')}")
+        print(f"  - Size: {ckpt2_path.stat().st_size / (1024 * 1024):.2f} MB")
+        
+        # So sánh val_loss
+        val_loss1 = ckpt1.get('val_loss')
+        val_loss2 = ckpt2.get('val_loss')
+        
+        if isinstance(val_loss1, float) and isinstance(val_loss2, float):
+            print(f"\n📊 So sánh:")
+            if val_loss1 < val_loss2:
+                print(f"  ✓ {checkpoint1_name} tốt hơn (val_loss thấp hơn {val_loss2 - val_loss1:.4f})")
+            elif val_loss2 < val_loss1:
+                print(f"  ✓ {checkpoint2_name} tốt hơn (val_loss thấp hơn {val_loss1 - val_loss2:.4f})")
+            else:
+                print(f"  = Cả 2 checkpoint có val_loss bằng nhau")
+        
+        # Kiểm tra finetune config
+        if 'finetune_config' in ckpt1:
+            print(f"\n🔧 {checkpoint1_name} là model đã finetune")
+        if 'finetune_config' in ckpt2:
+            print(f"\n🔧 {checkpoint2_name} là model đã finetune")
+        
+    except Exception as e:
+        print(f"❌ Lỗi khi so sánh: {e}")
+    
+    print("="*70)
+
+def main():
+    parser = argparse.ArgumentParser(description='Checkpoint Manager')
+    subparsers = parser.add_subparsers(dest='command', help='Command to execute')
+    
+    # List command
+    list_parser = subparsers.add_parser('list', help='List all checkpoints')
+    
+    # Switch command
+    switch_parser = subparsers.add_parser('switch', help='Switch checkpoint (copy to best_model.pt)')
+    switch_parser.add_argument('source', type=str, help='Source checkpoint name')
+    switch_parser.add_argument('--target', type=str, default='best_model.pt',
+                              help='Target checkpoint name (default: best_model.pt)')
+    switch_parser.add_argument('--no-backup', action='store_true',
+                              help='Do not backup old checkpoint')
+    
+    # Compare command
+    compare_parser = subparsers.add_parser('compare', help='Compare two checkpoints')
+    compare_parser.add_argument('checkpoint1', type=str, help='First checkpoint name')
+    compare_parser.add_argument('checkpoint2', type=str, help='Second checkpoint name')
+    
+    args = parser.parse_args()
+    
+    if args.command == 'list':
+        list_checkpoints()
+    elif args.command == 'switch':
+        switch_checkpoint(args.source, args.target, backup=not args.no_backup)
+    elif args.command == 'compare':
+        compare_checkpoints(args.checkpoint1, args.checkpoint2)
+    else:
+        parser.print_help()
+
+if __name__ == "__main__":
+    main()
+

src/complete_transformer.py

@@ -0,0 +1,539 @@
+"""
+COMPLETE TRANSFORMER MODEL
+Mô hình Transformer hoàn chỉnh cho dịch máy Seq2Seq
+"""
+
+import torch
+import torch.nn as nn
+from .transformer_encoder_decoder import (
+    Encoder, Decoder,
+    create_padding_mask, create_target_mask
+)
+
+# ============================================================================
+# TRANSFORMER MODEL
+# ============================================================================
+
+class Transformer(nn.Module):
+    """
+    Mô hình Transformer hoàn chỉnh cho Neural Machine Translation
+    
+    Args:
+        src_vocab_size: Kích thước vocabulary source language
+        tgt_vocab_size: Kích thước vocabulary target language
+        d_model: Dimension của model (mặc định 512)
+        n_layers: Số lượng encoder/decoder layers (mặc định 6)
+        n_heads: Số lượng attention heads (mặc định 8)
+        d_ff: Dimension của feed-forward network (mặc định 2048)
+        dropout: Dropout rate (mặc định 0.1)
+        max_len: Maximum sequence length (mặc định 5000)
+        pad_idx: Index của padding token (mặc định 0)
+    """
+    def __init__(
+        self,
+        src_vocab_size,
+        tgt_vocab_size,
+        d_model=512,
+        n_layers=6,
+        n_heads=8,
+        d_ff=2048,
+        dropout=0.1,
+        max_len=5000,
+        pad_idx=0
+    ):
+        super().__init__()
+        
+        self.pad_idx = pad_idx
+        
+        # Encoder
+        self.encoder = Encoder(
+            vocab_size=src_vocab_size,
+            d_model=d_model,
+            n_layers=n_layers,
+            n_heads=n_heads,
+            d_ff=d_ff,
+            dropout=dropout,
+            max_len=max_len
+        )
+        
+        # Decoder
+        self.decoder = Decoder(
+            vocab_size=tgt_vocab_size,
+            d_model=d_model,
+            n_layers=n_layers,
+            n_heads=n_heads,
+            d_ff=d_ff,
+            dropout=dropout,
+            max_len=max_len
+        )
+        
+        # Khởi tạo weights
+        self._init_weights()
+        
+    def _init_weights(self):
+        """
+        Khởi tạo weights theo Xavier Uniform
+        """
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+    
+    def forward(self, src, tgt):
+        """
+        Forward pass
+        
+        Args:
+            src: Source sequence [batch_size, src_len]
+            tgt: Target sequence [batch_size, tgt_len]
+            
+        Returns:
+            output: Logits [batch_size, tgt_len, tgt_vocab_size]
+        """
+        # Tạo masks
+        src_mask = create_padding_mask(src, self.pad_idx)
+        tgt_mask = create_target_mask(tgt, self.pad_idx)
+        
+        # Encoder
+        encoder_output = self.encoder(src, src_mask)
+        
+        # Decoder
+        output = self.decoder(tgt, encoder_output, src_mask, tgt_mask)
+        
+        return output
+    
+    def encode(self, src):
+        """
+        Chỉ chạy encoder (dùng khi inference)
+        
+        Args:
+            src: Source sequence [batch_size, src_len]
+            
+        Returns:
+            encoder_output: [batch_size, src_len, d_model]
+            src_mask: [batch_size, 1, 1, src_len]
+        """
+        src_mask = create_padding_mask(src, self.pad_idx)
+        encoder_output = self.encoder(src, src_mask)
+        return encoder_output, src_mask
+    
+    def decode(self, tgt, encoder_output, src_mask):
+        """
+        Chỉ chạy decoder (dùng khi inference)
+        
+        Args:
+            tgt: Target sequence [batch_size, tgt_len]
+            encoder_output: Encoder output [batch_size, src_len, d_model]
+            src_mask: Source mask [batch_size, 1, 1, src_len]
+            
+        Returns:
+            output: Logits [batch_size, tgt_len, tgt_vocab_size]
+        """
+        tgt_mask = create_target_mask(tgt, self.pad_idx)
+        output = self.decoder(tgt, encoder_output, src_mask, tgt_mask)
+        return output
+
+# ============================================================================
+# TRANSFORMER WITH SHARED VOCABULARY & WEIGHT TYING
+# ============================================================================
+
+class TransformerShared(nn.Module):
+    """
+    Transformer với Shared Vocabulary và Weight Tying
+    
+    Đặc điểm:
+    - Dùng chung 1 vocabulary cho cả source và target
+    - Embedding input và output layer chia sẻ weights (Weight Tying)
+    - Tiết kiệm ~50% parameters so với model riêng biệt
+    - Học được mối liên hệ trực tiếp giữa 2 ngôn ngữ tốt hơn
+    
+    Args:
+        vocab_size: Kích thước shared vocabulary
+        d_model: Dimension của model (mặc định 512)
+        n_layers: Số lượng encoder/decoder layers (mặc định 6)
+        n_heads: Số lượng attention heads (mặc định 8)
+        d_ff: Dimension của feed-forward network (mặc định 2048)
+        dropout: Dropout rate (mặc định 0.1)
+        max_len: Maximum sequence length (mặc định 5000)
+        pad_idx: Index của padding token (mặc định 0)
+        use_weight_tying: Có dùng weight tying không (mặc định True)
+    """
+    def __init__(
+        self,
+        vocab_size,
+        d_model=512,
+        n_layers=6,
+        n_heads=8,
+        d_ff=2048,
+        dropout=0.1,
+        max_len=5000,
+        pad_idx=0,
+        use_weight_tying=True
+    ):
+        super().__init__()
+        
+        self.pad_idx = pad_idx
+        self.d_model = d_model
+        self.use_weight_tying = use_weight_tying
+        
+        # --- KHÁC BIỆT LỚN NHẤT ---
+        # Chỉ tạo 1 Embedding matrix dùng cho cả 2 ngôn ngữ
+        from .transformer_components import Embedding, PositionalEncoding
+        self.shared_embedding = Embedding(vocab_size, d_model)
+        
+        # Positional Encoding
+        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)
+        
+        # Encoder (dùng shared embedding)
+        # Tạo Encoder nhưng sẽ thay embedding sau
+        self.encoder = Encoder(
+            vocab_size=vocab_size,  # Dùng chung vocab_size
+            d_model=d_model,
+            n_layers=n_layers,
+            n_heads=n_heads,
+            d_ff=d_ff,
+            dropout=dropout,
+            max_len=max_len
+        )
+        
+        # Thay thế embedding của encoder bằng shared embedding
+        # Quan trọng: Phải thay thế sau khi tạo Encoder
+        self.encoder.embedding = self.shared_embedding
+        
+        # Decoder (dùng shared embedding)
+        # Tạo Decoder nhưng sẽ thay embedding sau
+        self.decoder = Decoder(
+            vocab_size=vocab_size,  # Dùng chung vocab_size
+            d_model=d_model,
+            n_layers=n_layers,
+            n_heads=n_heads,
+            d_ff=d_ff,
+            dropout=dropout,
+            max_len=max_len
+        )
+        
+        # Thay thế embedding của decoder bằng shared embedding
+        # Quan trọng: Phải thay thế sau khi tạo Decoder
+        self.decoder.embedding = self.shared_embedding
+        
+        # QUAN TRỌNG: Bỏ qua fc_out của Decoder vì chúng ta dùng output_layer riêng
+        # Decoder.fc_out sẽ được thay thế bằng identity function
+        self.decoder.fc_out = nn.Identity()
+        
+        # Output layer
+        self.output_layer = nn.Linear(d_model, vocab_size, bias=False)
+        
+        # Khởi tạo weights TRƯỚC khi weight tying
+        self._init_weights()
+        
+        # --- KÍCH HOẠT WEIGHT TYING ---
+        if use_weight_tying:
+            # Dòng này giúp tiết kiệm 50% tham số vocab
+            # Embedding weight và output layer weight chia sẻ nhau
+            # QUAN TRỌNG: Phải gán SAU _init_weights() để không bị reset
+            self.output_layer.weight = self.shared_embedding.embedding.weight
+            print("✓ Weight Tying enabled: Embedding và Output layer chia sẻ weights")
+        
+    def _init_weights(self):
+        """
+        Khởi tạo weights theo Xavier Uniform
+        """
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+    
+    def forward(self, src, tgt):
+        """
+        Forward pass
+        
+        Args:
+            src: Source sequence [batch_size, src_len]
+            tgt: Target sequence [batch_size, tgt_len]
+            
+        Returns:
+            output: Logits [batch_size, tgt_len, vocab_size]
+        """
+        # Tạo masks
+        src_mask = create_padding_mask(src, self.pad_idx)
+        tgt_mask = create_target_mask(tgt, self.pad_idx)
+        
+        # Encoder (dùng shared embedding)
+        encoder_output = self.encoder(src, src_mask)
+        
+        # Decoder (dùng shared embedding)
+        decoder_output = self.decoder(tgt, encoder_output, src_mask, tgt_mask)
+        
+        # Debug: Kiểm tra shape
+        # print(f"DEBUG: decoder_output shape: {decoder_output.shape}")
+        # print(f"DEBUG: output_layer weight shape: {self.output_layer.weight.shape}")
+        
+        # Output layer (có thể share weight với embedding nếu use_weight_tying=True)
+        output = self.output_layer(decoder_output)
+        
+        return output
+    
+    def encode(self, src):
+        """
+        Chỉ chạy encoder (dùng khi inference)
+        
+        Args:
+            src: Source sequence [batch_size, src_len]
+            
+        Returns:
+            encoder_output: [batch_size, src_len, d_model]
+            src_mask: [batch_size, 1, 1, src_len]
+        """
+        src_mask = create_padding_mask(src, self.pad_idx)
+        encoder_output = self.encoder(src, src_mask)
+        return encoder_output, src_mask
+    
+    def decode(self, tgt, encoder_output, src_mask):
+        """
+        Chỉ chạy decoder (dùng khi inference)
+        
+        Args:
+            tgt: Target sequence [batch_size, tgt_len]
+            encoder_output: Encoder output [batch_size, src_len, d_model]
+            src_mask: Source mask [batch_size, 1, 1, src_len]
+            
+        Returns:
+            output: Logits [batch_size, tgt_len, vocab_size]
+        """
+        tgt_mask = create_target_mask(tgt, self.pad_idx)
+        decoder_output = self.decoder(tgt, encoder_output, src_mask, tgt_mask)
+        output = self.output_layer(decoder_output)
+        return output
+
+# ============================================================================
+# MODEL CONFIGURATION
+# ============================================================================
+
+def get_model_config(model_size='base'):
+    """
+    Trả về config cho các kích thước model khác nhau
+    
+    Args:
+        model_size: 'tiny', 'small', 'base', 'large'
+        
+    Returns:
+        config: Dictionary chứa hyperparameters
+    """
+    configs = {
+        'tiny': {
+            'd_model': 256,
+            'n_layers': 2,
+            'n_heads': 4,
+            'd_ff': 1024,
+            'dropout': 0.1
+        },
+        'small': {
+            'd_model': 256,
+            'n_layers': 4,
+            'n_heads': 8,
+            'd_ff': 1024,
+            'dropout': 0.1
+        },
+        'medium': {  # ~25M parameters với 32k vocab + weight tying
+            'd_model': 384,
+            'n_layers': 5,
+            'n_heads': 8,
+            'd_ff': 1536,
+            'dropout': 0.1
+        },
+        'custom_25m': {  # ~25M parameters với 32k vocab + weight tying
+            'd_model': 384,
+            'n_layers': 6,
+            'n_heads': 8,
+            'd_ff': 1536,  # 4 * d_model
+            'dropout': 0.1
+        },
+        'base': {
+            'd_model': 512,
+            'n_layers': 6,
+            'n_heads': 8,
+            'd_ff': 2048,
+            'dropout': 0.1
+        },
+        'large': {
+            'd_model': 1024,
+            'n_layers': 6,
+            'n_heads': 16,
+            'd_ff': 4096,
+            'dropout': 0.1
+        }
+    }
+    
+    return configs.get(model_size, configs['base'])
+
+def create_model(src_vocab_size, tgt_vocab_size, model_size='base', pad_idx=0, 
+                use_shared_vocab=True, use_weight_tying=True):
+    """
+    Tạo Transformer model với Shared Vocabulary
+    
+    Args:
+        src_vocab_size: Kích thước shared vocabulary
+        tgt_vocab_size: Bỏ qua (giữ để tương thích, phải = src_vocab_size)
+        model_size: Kích thước model ('tiny', 'small', 'base', 'large')
+        pad_idx: Padding index
+        use_shared_vocab: Luôn True (giữ để tương thích)
+        use_weight_tying: Có dùng weight tying không (mặc định True)
+        
+    Returns:
+        model: TransformerShared model
+        config: Model configuration
+    """
+    config = get_model_config(model_size)
+    
+    # Luôn dùng shared vocabulary
+    vocab_size = src_vocab_size
+    model = TransformerShared(
+        vocab_size=vocab_size,
+        d_model=config['d_model'],
+        n_layers=config['n_layers'],
+        n_heads=config['n_heads'],
+        d_ff=config['d_ff'],
+        dropout=config['dropout'],
+        pad_idx=pad_idx,
+        use_weight_tying=use_weight_tying
+    )
+    
+    return model, config
+
+# ============================================================================
+# UTILITY FUNCTIONS
+# ============================================================================
+
+def count_parameters(model):
+    """
+    Đếm số lượng parameters của model
+    
+    Args:
+        model: PyTorch model
+        
+    Returns:
+        total: Tổng số parameters
+        trainable: Số parameters có thể train
+    """
+    total = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    
+    return total, trainable
+
+def print_model_info(model, model_size='base', use_shared_vocab=False):
+    """
+    In thông tin về model
+    
+    Args:
+        model: Transformer model
+        model_size: Kích thước model
+        use_shared_vocab: Có dùng shared vocabulary không
+    """
+    total_params, trainable_params = count_parameters(model)
+    
+    print("="*70)
+    print("THÔNG TIN MÔ HÌNH TRANSFORMER")
+    print("="*70)
+    print(f"\nKích thước model: {model_size.upper()}")
+    
+    if use_shared_vocab:
+        print(f"  Mode: SHARED VOCABULARY + WEIGHT TYING")
+        if isinstance(model, TransformerShared) and model.use_weight_tying:
+            print(f"  ✓ Weight Tying: Enabled (tiết kiệm ~50% vocab params)")
+    else:
+        print(f"  Mode: SEPARATE VOCABULARIES")
+    
+    print(f"\nSố lượng parameters:")
+    print(f"  - Total: {total_params:,}")
+    print(f"  - Trainable: {trainable_params:,}")
+    print(f"  - Model size: ~{total_params * 4 / (1024**2):.2f} MB (float32)")
+    
+    config = get_model_config(model_size)
+    print(f"\nCấu hình:")
+    print(f"  - d_model: {config['d_model']}")
+    print(f"  - n_layers: {config['n_layers']}")
+    print(f"  - n_heads: {config['n_heads']}")
+    print(f"  - d_ff: {config['d_ff']}")
+    print(f"  - dropout: {config['dropout']}")
+    print("="*70)
+
+# ============================================================================
+# TEST COMPLETE MODEL
+# ============================================================================
+
+if __name__ == "__main__":
+    print("="*70)
+    print("KIỂM TRA TRANSFORMER MODEL HOÀN CHỈNH")
+    print("="*70)
+    
+    # Hyperparameters
+    src_vocab_size = 10000
+    tgt_vocab_size = 8000
+    batch_size = 4
+    src_len = 15
+    tgt_len = 20
+    pad_idx = 0
+    
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"\nDevice: {device}\n")
+    
+    # Test với các kích thước model khác nhau
+    for model_size in ['tiny', 'small', 'base']:
+        print(f"\n{'='*70}")
+        print(f"TEST MODEL SIZE: {model_size.upper()}")
+        print(f"{'='*70}\n")
+        
+        # Tạo model
+        model, config = create_model(
+            src_vocab_size=src_vocab_size,
+            tgt_vocab_size=tgt_vocab_size,
+            model_size=model_size,
+            pad_idx=pad_idx
+        )
+        model = model.to(device)
+        
+        # In thông tin model
+        print_model_info(model, model_size)
+        
+        # Tạo dummy data
+        src = torch.randint(1, src_vocab_size, (batch_size, src_len)).to(device)
+        tgt = torch.randint(1, tgt_vocab_size, (batch_size, tgt_len)).to(device)
+        
+        # Forward pass
+        print(f"\nForward pass:")
+        print(f"  Source shape: {src.shape}")
+        print(f"  Target shape: {tgt.shape}")
+        
+        with torch.no_grad():
+            output = model(src, tgt)
+        
+        print(f"  Output shape: {output.shape}")
+        print(f"  Expected: [{batch_size}, {tgt_len}, {tgt_vocab_size}]")
+        print(f"  ✓ Shape correct!")
+        
+        # Test encode và decode riêng
+        print(f"\nTest encode & decode separately:")
+        with torch.no_grad():
+            encoder_output, src_mask = model.encode(src)
+            decoder_output = model.decode(tgt, encoder_output, src_mask)
+        
+        print(f"  Encoder output shape: {encoder_output.shape}")
+        print(f"  Decoder output shape: {decoder_output.shape}")
+        print(f"  ✓ Encode/Decode work correctly!")
+        
+        # Kiểm tra output giống nhau
+        print(f"\nVerify output consistency:")
+        with torch.no_grad():
+            output_combined = model(src, tgt)
+        
+        is_same = torch.allclose(output_combined, decoder_output, atol=1e-6)
+        print(f"  Forward == Encode+Decode: {is_same}")
+        print(f"  ✓ Model is consistent!")
+    
+    print("\n" + "="*70)
+    print("✓ TẤT CẢ TESTS PASSED!")
+    print("="*70)
+    
+    print("\n📝 GỢI Ý SỬ DỤNG:")
+    print("  - Dùng 'tiny' để debug và test nhanh")
+    print("  - Dùng 'small' để train trên CPU hoặc GPU nhỏ")
+    print("  - Dùng 'base' để có kết quả tốt (cần GPU)")
+    print("  - Dùng 'large' chỉ khi có GPU mạnh")

src/config.py

@@ -0,0 +1,36 @@
+"""
+CONFIGURATION FILE
+Centralized configuration for the project
+"""
+
+import torch
+from pathlib import Path
+
+class Config:
+    """Global configuration"""
+    
+    # Paths
+    PROJECT_ROOT = Path(__file__).parent.parent
+    DATA_DIR = PROJECT_ROOT / 'data'
+    CHECKPOINT_DIR = PROJECT_ROOT / 'checkpoints'
+    RESULTS_DIR = PROJECT_ROOT / 'results'
+    LOGS_DIR = PROJECT_ROOT / 'logs'
+    
+    # Data
+    PROCESSED_DATA_DIR = DATA_DIR / 'processed'
+    VI_VOCAB_PATH = PROCESSED_DATA_DIR / 'vi_vocab.pkl'
+    EN_VOCAB_PATH = PROCESSED_DATA_DIR / 'en_vocab.pkl'
+    PROCESSED_DATA_PATH = PROCESSED_DATA_DIR / 'processed_data.pkl'
+    
+    # Model
+    MODEL_SIZE = 'base'  # 'tiny', 'small', 'base', 'large'
+    PAD_IDX = 0
+    
+    # Training
+    NUM_EPOCHS = 20
+    BATCH_SIZE = 32
+    WARMUP_STEPS = 4000
+    LABEL_SMOOTHING = 0.1
+    
+    # Device
+    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

src/create_bidirectional_dataset.py

@@ -0,0 +1,91 @@
+"""
+Tạo Bidirectional Dataset từ dataset hiện tại
+Reverse các cặp câu để model học được cả hai chiều
+"""
+
+import pickle
+from pathlib import Path
+from tqdm import tqdm
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+PROCESSED_DATA_DIR = PROJECT_ROOT / 'data' / 'processed'
+
+def create_bidirectional_dataset():
+    """
+    Tạo bidirectional dataset bằng cách:
+    1. Giữ nguyên các cặp vi→en
+    2. Thêm các cặp en→vi (reverse)
+    """
+    print("="*70)
+    print("TẠO BIDIRECTIONAL DATASET")
+    print("="*70)
+    
+    # Load data hiện tại
+    data_path = PROCESSED_DATA_DIR / 'processed_data_shared.pkl'
+    if not data_path.exists():
+        print(f"❌ Không tìm thấy: {data_path}")
+        print("   Vui lòng chạy: python src/2_encode_data.py trước")
+        return
+    
+    print(f"\n📂 Loading dataset từ: {data_path}")
+    with open(data_path, 'rb') as f:
+        data = pickle.load(f)
+    
+    # Hiển thị thống kê ban đầu
+    print("\n📊 Thống kê dataset hiện tại:")
+    for split in ['train', 'validation', 'test']:
+        if split in data:
+            print(f"  {split}: {len(data[split])} cặp câu (vi→en)")
+    
+    # Tạo bidirectional data
+    print("\n🔄 Tạo bidirectional dataset...")
+    bidirectional_data = {
+        'train': [],
+        'validation': [],
+        'test': []
+    }
+    
+    for split in ['train', 'validation', 'test']:
+        if split not in data:
+            continue
+            
+        print(f"\n  Processing {split}...")
+        original_count = len(data[split])
+        
+        # Giữ nguyên chiều vi→en
+        bidirectional_data[split].extend(data[split])
+        
+        # Thêm chiều ngược lại en→vi
+        for src_ids, tgt_ids in tqdm(data[split], desc=f"  Reversing {split}"):
+            bidirectional_data[split].append((tgt_ids, src_ids))  # Swap!
+        
+        new_count = len(bidirectional_data[split])
+        print(f"  ✓ {split}: {original_count} → {new_count} cặp câu (x2)")
+    
+    # Lưu lại
+    output_path = PROCESSED_DATA_DIR / 'processed_data_bidirectional.pkl'
+    print(f"\n💾 Lưu bidirectional dataset vào: {output_path}")
+    with open(output_path, 'wb') as f:
+        pickle.dump(bidirectional_data, f)
+    
+    # Thống kê cuối cùng
+    print("\n" + "="*70)
+    print("KẾT QUẢ")
+    print("="*70)
+    total_pairs = 0
+    for split in ['train', 'validation', 'test']:
+        count = len(bidirectional_data[split])
+        total_pairs += count
+        print(f"  {split}: {count:,} cặp câu")
+    print(f"\n  Tổng cộng: {total_pairs:,} cặp câu")
+    print(f"  ✓ Dataset đã sẵn sàng để train bidirectional!")
+    print("="*70)
+    
+    print("\n📝 Bước tiếp theo:")
+    print("  1. Cập nhật training code để dùng 'processed_data_bidirectional.pkl'")
+    print("  2. Train lại model với dataset mới")
+    print("  3. Model sẽ học được cả vi→en và en→vi")
+
+if __name__ == '__main__':
+    create_bidirectional_dataset()
+

src/data_preprocessing.py

@@ -0,0 +1,468 @@
+"""
+DATA PREPROCESSING - PHIÊN BẢN ĐƠN GIẢN
+Load dữ liệu từ nhiều nguồn
+"""
+
+import re
+import json
+import pickle
+import os
+from collections import Counter
+from typing import List, Dict
+
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from pathlib import Path
+
+# ============================================================================
+# PATH CONFIGURATION
+# ============================================================================
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DATA_DIR = PROJECT_ROOT / 'data'
+RAW_DATA_DIR = DATA_DIR / 'raw'
+PROCESSED_DATA_DIR = DATA_DIR / 'processed'
+CUSTOM_CSV_PATH = PROJECT_ROOT / 'mtet_cleaned.csv'
+# Set to 0 or negative to use all rows, or a positive number to limit
+# Default: 0 (use all rows) - change to a number like '500000' to limit
+_max_rows_env = os.environ.get('MTET_MAX_ROWS', '0')
+CUSTOM_CSV_MAX_ROWS = int(_max_rows_env) if _max_rows_env and int(_max_rows_env) > 0 else None
+
+# ============================================================================
+# 1. LOAD DATASET
+# ============================================================================
+
+def load_iwslt_data():
+    """
+    Load IWSLT dataset với nhiều fallback options
+    """
+    # Option 0: Custom CSV (e.g., MTet mirror)
+    if CUSTOM_CSV_PATH.exists():
+        print("Đang tải dataset từ custom CSV:", CUSTOM_CSV_PATH)
+        try:
+            return load_custom_csv(CUSTOM_CSV_PATH)
+        except Exception as e:  # noqa: BLE001
+            print(f"  ✗ Failed custom CSV: {e}")
+
+    print("Đang tải dataset IWSLT Vi-En...")
+    
+    # Option 1: Thử load từ Hugging Face
+    try:
+        from datasets import load_dataset
+        print("  Thử Option 1: Hugging Face mt_eng_vietnamese...")
+        dataset = load_dataset("mt_eng_vietnamese", "iwslt2015-vi-en", trust_remote_code=True)
+        print("  ✓ Loaded successfully!")
+        return dataset
+    except Exception as e:
+        print(f"  ✗ Failed: {e}")
+    
+    # Option 2: Thử dataset khác
+    try:
+        from datasets import load_dataset
+        print("  Thử Option 2: Hugging Face iwslt2017...")
+        dataset = load_dataset("iwslt2017", "iwslt2017-vi-en", trust_remote_code=True)
+        print("  ✓ Loaded successfully!")
+        return dataset
+    except Exception as e:
+        print(f"  ✗ Failed: {e}")
+    
+    # Option 3: Load từ local files (nếu đã download)
+    try:
+        print("  Thử Option 3: Load từ local files...")
+        dataset = load_from_local()
+        print("  ✓ Loaded successfully!")
+        return dataset
+    except Exception as e:
+        print(f"  ✗ Failed: {e}")
+    
+    # Option 4: Tạo sample dataset nhỏ
+    print("  ⚠️  Dùng Option 4: Sample dataset (chỉ để test)")
+    print("  📝 Tải dataset thật từ: https://nlp.stanford.edu/projects/nmt/")
+    return create_sample_dataset()
+
+def load_from_local():
+    """
+    Load dataset từ file local
+    Yêu cầu: Đã download và đặt trong data/raw/
+    """
+    from datasets import Dataset, DatasetDict
+    
+    data_dir = RAW_DATA_DIR
+    
+    # Đọc các file
+    splits = {}
+    for split in ['train', 'dev', 'test']:
+        vi_file = data_dir / f'{split}.vi'
+        en_file = data_dir / f'{split}.en'
+        
+        if not vi_file.exists() or not en_file.exists():
+            raise FileNotFoundError(f"Không tìm thấy {vi_file} hoặc {en_file}")
+        
+        with open(vi_file, 'r', encoding='utf-8') as f:
+            vi_lines = f.readlines()
+        
+        with open(en_file, 'r', encoding='utf-8') as f:
+            en_lines = f.readlines()
+        
+        # Tạo dataset
+        translations = [
+            {'vi': vi.strip(), 'en': en.strip()}
+            for vi, en in zip(vi_lines, en_lines)
+        ]
+        
+        splits[split if split != 'dev' else 'validation'] = Dataset.from_dict({
+            'translation': translations
+        })
+    
+    return DatasetDict(splits)
+
+
+def load_custom_csv(csv_path, train_ratio=0.98, val_ratio=0.01, seed=42, max_rows=None):
+    """
+    Load dataset từ CSV custom (ví dụ MTet) với cột src/tgt hoặc vi/en.
+    """
+    from datasets import Dataset, DatasetDict
+
+    df = pd.read_csv(csv_path)
+    if {'src', 'tgt'}.issubset(df.columns):
+        df = df.rename(columns={'src': 'vi', 'tgt': 'en'})
+    elif {'vi', 'en'}.issubset(df.columns):
+        pass
+    else:
+        raise ValueError("CSV cần có cột 'src'/'tgt' hoặc 'vi'/'en'.")
+
+    df = df[['vi', 'en']].dropna().reset_index(drop=True)
+    if df.empty:
+        raise ValueError("CSV không có dữ liệu hợp lệ.")
+
+    original_total = len(df)
+    max_rows = max_rows if max_rows is not None else CUSTOM_CSV_MAX_ROWS
+    if max_rows is not None and max_rows > 0 and original_total > max_rows:
+        print(f"  🔎 Sampling {max_rows} rows (từ {original_total}) để phù hợp bộ nhớ.")
+        df = df.sample(n=max_rows, random_state=seed).reset_index(drop=True)
+    elif max_rows is None or max_rows <= 0:
+        print(f"  📊 Using all {original_total} rows from CSV.")
+
+    df = df.sample(frac=1.0, random_state=seed).reset_index(drop=True)
+    total = len(df)
+    train_end = max(1, int(total * train_ratio))
+    val_end = train_end + max(1, int(total * val_ratio))
+    if val_end >= total:
+        val_end = total - 1
+    if val_end <= train_end:
+        val_end = min(total - 1, train_end + 1)
+
+    splits = {
+        'train': df.iloc[:train_end],
+        'validation': df.iloc[train_end:val_end],
+        'test': df.iloc[val_end:]
+    }
+
+    dataset_dict = {}
+    for split, split_df in splits.items():
+        examples = [{'vi': row['vi'], 'en': row['en']} for _, row in split_df.iterrows()]
+        dataset_dict[split] = Dataset.from_dict({'translation': examples})
+
+    hf_dataset = DatasetDict(dataset_dict)
+    print(f"  📊 Custom CSV dataset: train={len(hf_dataset['train'])}, "
+          f"val={len(hf_dataset['validation'])}, test={len(hf_dataset['test'])}")
+
+    return hf_dataset
+def create_sample_dataset():
+    """
+    Tạo dataset mẫu để test code
+    """
+    from datasets import Dataset, DatasetDict
+    
+    # Sample Vietnamese-English pairs
+    samples = [
+        ('xin chào', 'hello'),
+        ('tạm biệt', 'goodbye'),
+        ('cảm ơn', 'thank you'),
+        ('xin lỗi', 'sorry'),
+        ('tôi tên là Nam', 'my name is Nam'),
+        ('tôi là sinh viên', 'i am a student'),
+        ('hôm nay thời tiết đẹp', 'the weather is nice today'),
+        ('tôi đang học tiếng Anh', 'i am learning English'),
+        ('bạn khỏe không', 'how are you'),
+        ('tôi yêu Việt Nam', 'i love Vietnam'),
+        ('chúc bạn một ngày tốt lành', 'have a nice day'),
+        ('rất vui được gặp bạn', 'nice to meet you'),
+        ('tôi đang làm việc', 'i am working'),
+        ('hãy giúp tôi', 'please help me'),
+        ('tôi không hiểu', 'i do not understand'),
+        ('bạn nói tiếng Anh không', 'do you speak English'),
+        ('tôi đói bụng', 'i am hungry'),
+        ('mấy giờ rồi', 'what time is it'),
+        ('tôi muốn đi', 'i want to go'),
+        ('đây là gì', 'what is this'),
+    ]
+    
+    # Tạo train set (lớn hơn)
+    train_data = []
+    for _ in range(200):  # Nhân lên
+        for vi, en in samples:
+            train_data.append({'vi': vi, 'en': en})
+    
+    # Val và test
+    val_data = [{'vi': vi, 'en': en} for vi, en in samples[:5]] * 10
+    test_data = [{'vi': vi, 'en': en} for vi, en in samples[5:10]] * 10
+    
+    dataset = DatasetDict({
+        'train': Dataset.from_dict({'translation': train_data}),
+        'validation': Dataset.from_dict({'translation': val_data}),
+        'test': Dataset.from_dict({'translation': test_data})
+    })
+    
+    print(f"  📊 Sample dataset: {len(train_data)} train, {len(val_data)} val, {len(test_data)} test")
+    
+    return dataset
+
+# ============================================================================
+# 2. CLEANING
+# ============================================================================
+
+def clean_text(text: str, lang: str = 'vi') -> str:
+    """Làm sạch văn bản"""
+    text = text.lower()
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    text = re.sub(r'([.,!?;:])', r' \1 ', text)
+    text = re.sub(r'\s+', ' ', text)
+    return text
+
+def clean_dataset(dataset):
+    """Làm sạch toàn bộ dataset"""
+    print("\n" + "="*70)
+    print("LÀM SẠCH DỮ LIỆU")
+    print("="*70)
+    
+    cleaned_data = {
+        'train': [],
+        'validation': [],
+        'test': []
+    }
+    
+    for split in ['train', 'validation', 'test']:
+        print(f"\nĐang làm sạch {split} set...")
+        
+        for example in tqdm(dataset[split], desc=f"Cleaning {split}"):
+            vi_text = clean_text(example['translation']['vi'], 'vi')
+            en_text = clean_text(example['translation']['en'], 'en')
+            
+            # Filter
+            vi_words = len(vi_text.split())
+            en_words = len(en_text.split())
+            
+            if 1 <= vi_words <= 100 and 1 <= en_words <= 100:
+                ratio = max(vi_words, en_words) / max(min(vi_words, en_words), 1)
+                if ratio <= 3:
+                    cleaned_data[split].append({
+                        'vi': vi_text,
+                        'en': en_text
+                    })
+        
+        print(f"  Giữ lại: {len(cleaned_data[split])}/{len(dataset[split])} cặp câu")
+    
+    return cleaned_data
+
+# ============================================================================
+# 3. VOCABULARY
+# ============================================================================
+
+class Vocabulary:
+    """Lớp xây dựng và quản lý từ điển"""
+    def __init__(self, freq_threshold=2):
+        self.freq_threshold = freq_threshold
+        self.word2idx = {}
+        self.idx2word = {}
+        self.word_freq = Counter()
+        
+        self.PAD_TOKEN = '<pad>'
+        self.SOS_TOKEN = '<sos>'
+        self.EOS_TOKEN = '<eos>'
+        self.UNK_TOKEN = '<unk>'
+        
+        self.PAD_IDX = 0
+        self.SOS_IDX = 1
+        self.EOS_IDX = 2
+        self.UNK_IDX = 3
+        
+    def build_vocabulary(self, sentences: List[str]):
+        """Xây dựng từ điển từ danh sách câu"""
+        print("Đang xây dựng từ điển...")
+        
+        for sentence in tqdm(sentences, desc="Counting words"):
+            for word in sentence.split():
+                self.word_freq[word] += 1
+        
+        idx = 0
+        for token in [self.PAD_TOKEN, self.SOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]:
+            self.word2idx[token] = idx
+            self.idx2word[idx] = token
+            idx += 1
+        
+        for word, freq in self.word_freq.most_common():
+            if freq >= self.freq_threshold:
+                self.word2idx[word] = idx
+                self.idx2word[idx] = word
+                idx += 1
+        
+        print(f"Kích thước từ điển: {len(self.word2idx)}")
+        
+    def encode(self, sentence: str) -> List[int]:
+        """Chuyển câu thành chuỗi index"""
+        tokens = [self.SOS_IDX]
+        for word in sentence.split():
+            tokens.append(self.word2idx.get(word, self.UNK_IDX))
+        tokens.append(self.EOS_IDX)
+        return tokens
+    
+    def decode(self, indices: List[int]) -> str:
+        """Chuyển chuỗi index thành câu"""
+        words = []
+        for idx in indices:
+            if idx == self.EOS_IDX:
+                break
+            if idx not in [self.PAD_IDX, self.SOS_IDX]:
+                words.append(self.idx2word.get(idx, self.UNK_TOKEN))
+        return ' '.join(words)
+    
+    def __len__(self):
+        return len(self.word2idx)
+
+# ============================================================================
+# 4. PREPARE DATA
+# ============================================================================
+
+def prepare_data_with_padding(cleaned_data, vi_vocab, en_vocab, max_len=100):
+    """Chuẩn bị dữ liệu với padding"""
+    print("\n" + "="*70)
+    print("CHUẨN BỊ DỮ LIỆU VỚI PADDING")
+    print("="*70)
+    
+    processed_data = {}
+    
+    for split in ['train', 'validation', 'test']:
+        print(f"\nĐang xử lý {split} set...")
+        
+        src_data = []
+        tgt_data = []
+        
+        for item in tqdm(cleaned_data[split], desc=f"Processing {split}"):
+            src_tokens = vi_vocab.encode(item['vi'])
+            tgt_tokens = en_vocab.encode(item['en'])
+            
+            if len(src_tokens) > max_len:
+                src_tokens = src_tokens[:max_len-1] + [vi_vocab.EOS_IDX]
+            if len(tgt_tokens) > max_len:
+                tgt_tokens = tgt_tokens[:max_len-1] + [en_vocab.EOS_IDX]
+            
+            src_data.append(src_tokens)
+            tgt_data.append(tgt_tokens)
+        
+        processed_data[split] = {
+            'src': src_data,
+            'tgt': tgt_data
+        }
+        
+        print(f"  Số mẫu: {len(src_data)}")
+    
+    return processed_data
+
+# ============================================================================
+# 5. SAVE
+# ============================================================================
+
+def save_data(cleaned_data, vi_vocab, en_vocab, processed_data):
+    """Lưu dữ liệu và từ điển"""
+    print("\n" + "="*70)
+    print("LƯU DỮ LIỆU")
+    print("="*70)
+    
+    os.makedirs(PROCESSED_DATA_DIR, exist_ok=True)
+    
+    with open(PROCESSED_DATA_DIR / 'cleaned_data.json', 'w', encoding='utf-8') as f:
+        json.dump(cleaned_data, f, ensure_ascii=False, indent=2)
+    print("✓ Đã lưu cleaned_data.json")
+    
+    with open(PROCESSED_DATA_DIR / 'vi_vocab.pkl', 'wb') as f:
+        pickle.dump(vi_vocab, f)
+    print("✓ Đã lưu vi_vocab.pkl")
+    
+    with open(PROCESSED_DATA_DIR / 'en_vocab.pkl', 'wb') as f:
+        pickle.dump(en_vocab, f)
+    print("✓ Đã lưu en_vocab.pkl")
+    
+    with open(PROCESSED_DATA_DIR / 'processed_data.pkl', 'wb') as f:
+        pickle.dump(processed_data, f)
+    print("✓ Đã lưu processed_data.pkl")
+
+# ============================================================================
+# 6. MAIN
+# ============================================================================
+
+def main():
+    """Pipeline xử lý dữ liệu hoàn chỉnh"""
+    
+    print("\n" + "="*70)
+    print("DATA PREPROCESSING PIPELINE")
+    print("="*70)
+    
+    # Load dataset
+    dataset = load_iwslt_data()
+    
+    # Explore
+    print("\n" + "="*70)
+    print("THỐNG KÊ DỮ LIỆU")
+    print("="*70)
+    for split in dataset.keys():
+        print(f"  - {split}: {len(dataset[split])} cặp câu")
+    
+    # Clean
+    cleaned_data = clean_dataset(dataset)
+    
+    # Build vocabularies
+    print("\n" + "="*70)
+    print("XÂY DỰNG TỪ ĐIỂN")
+    print("="*70)
+    
+    vi_sentences = [item['vi'] for item in cleaned_data['train']]
+    en_sentences = [item['en'] for item in cleaned_data['train']]
+    
+    print("\nTiếng Việt:")
+    vi_vocab = Vocabulary(freq_threshold=2)
+    vi_vocab.build_vocabulary(vi_sentences)
+    
+    print("\nTiếng Anh:")
+    en_vocab = Vocabulary(freq_threshold=2)
+    en_vocab.build_vocabulary(en_sentences)
+    
+    # Prepare data
+    processed_data = prepare_data_with_padding(cleaned_data, vi_vocab, en_vocab)
+    
+    # Save
+    save_data(cleaned_data, vi_vocab, en_vocab, processed_data)
+    
+    print("\n" + "="*70)
+    print("✓✓✓ HOÀN TẤT XỬ LÝ DỮ LIỆU! ✓✓✓")
+    print("="*70)
+    print(f"\nTóm tắt:")
+    print(f"  - Kích thước từ điển VI: {len(vi_vocab)}")
+    print(f"  - Kích thước từ điển EN: {len(en_vocab)}")
+    print(f"  - Train samples: {len(processed_data['train']['src'])}")
+    print(f"  - Validation samples: {len(processed_data['validation']['src'])}")
+    print(f"  - Test samples: {len(processed_data['test']['src'])}")
+    
+    print("\n📝 LƯU Ý:")
+    if len(processed_data['train']['src']) < 10000:
+        print("  ⚠️  Đang dùng sample dataset nhỏ!")
+        print("  📥 Để có BLEU tốt, hãy download dataset thật:")
+        print("     1. Vào: https://nlp.stanford.edu/projects/nmt/")
+        print("     2. Download: IWSLT'15 Vietnamese-English")
+        print("     3. Đặt vào: data/raw/")
+        print("     4. Chạy lại script này")
+
+if __name__ == "__main__":
+    main()

src/dataloader_module.py

@@ -0,0 +1,420 @@
+"""
+DATALOADER MODULE - TỐI ỨU HÓA BATCH PROCESSING (FIXED FOR WINDOWS)
+Xử lý batch với dynamic padding để tăng tốc huấn luyện
+"""
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+from torch.nn.utils.rnn import pad_sequence
+import pickle
+import numpy as np
+from pathlib import Path
+
+# ============================================================================
+# PATH CONFIGURATION
+# ============================================================================
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+PROCESSED_DATA_DIR = PROJECT_ROOT / 'data' / 'processed'
+
+# ============================================================================
+# 1. CUSTOM DATASET
+# ============================================================================
+
+class TranslationDataset(Dataset):
+    """
+    Dataset class cho dữ liệu dịch máy
+    """
+    def __init__(self, src_data, tgt_data):
+        """
+        Args:
+            src_data: List of source sequences (đã encode thành indices)
+            tgt_data: List of target sequences (đã encode thành indices)
+        """
+        assert len(src_data) == len(tgt_data), "Source và Target phải có cùng số lượng mẫu"
+        
+        self.src_data = src_data
+        self.tgt_data = tgt_data
+        
+    def __len__(self):
+        return len(self.src_data)
+    
+    def __getitem__(self, idx):
+        """
+        Trả về một cặp (source, target) dưới dạng tensor
+        """
+        src = torch.LongTensor(self.src_data[idx])
+        tgt = torch.LongTensor(self.tgt_data[idx])
+        
+        return src, tgt
+
+# ============================================================================
+# 2. COLLATE FUNCTION - DYNAMIC PADDING
+# ============================================================================
+
+def collate_fn(batch, pad_idx=0):
+    """
+    Collate function với dynamic padding
+    Chỉ pad đến độ dài câu dài nhất trong batch, không phải max_len cố định
+    
+    Args:
+        batch: List of (src, tgt) tuples
+        pad_idx: Index của padding token
+        
+    Returns:
+        src_batch: Padded source sequences [batch_size, max_src_len]
+        tgt_batch: Padded target sequences [batch_size, max_tgt_len]
+        src_lengths: Độ dài thực của mỗi source sequence
+        tgt_lengths: Độ dài thực của mỗi target sequence
+    """
+    # Tách source và target
+    src_batch, tgt_batch = zip(*batch)
+    
+    # Lấy độ dài thực của mỗi sequence (trước khi pad)
+    src_lengths = torch.LongTensor([len(s) for s in src_batch])
+    tgt_lengths = torch.LongTensor([len(t) for t in tgt_batch])
+    
+    # Padding - pad_sequence tự động pad đến độ dài max trong batch
+    src_batch = pad_sequence(src_batch, batch_first=True, padding_value=pad_idx)
+    tgt_batch = pad_sequence(tgt_batch, batch_first=True, padding_value=pad_idx)
+    
+    return src_batch, tgt_batch, src_lengths, tgt_lengths
+
+# ============================================================================
+# COLLATE WRAPPER - FIX CHO WINDOWS MULTIPROCESSING
+# ============================================================================
+
+class CollateWrapper:
+    """
+    Wrapper cho collate_fn để tránh lỗi pickle với lambda trên Windows
+    """
+    def __init__(self, pad_idx=0):
+        self.pad_idx = pad_idx
+    
+    def __call__(self, batch):
+        return collate_fn(batch, self.pad_idx)
+
+# ============================================================================
+# 3. TẠO DATALOADER
+# ============================================================================
+
+def create_dataloaders(processed_data, batch_size=32, num_workers=0):
+    """
+    Tạo DataLoader cho train, validation và test
+    
+    Args:
+        processed_data: Dict chứa 'train', 'validation', 'test'
+        batch_size: Kích thước batch
+        num_workers: Số worker threads (set = 0 cho Windows)
+        
+    Returns:
+        train_loader, val_loader, test_loader
+    """
+    print("\n" + "="*70)
+    print("TẠO DATALOADERS")
+    print("="*70)
+    
+    # Tạo datasets
+    train_dataset = TranslationDataset(
+        processed_data['train']['src'],
+        processed_data['train']['tgt']
+    )
+    
+    val_dataset = TranslationDataset(
+        processed_data['validation']['src'],
+        processed_data['validation']['tgt']
+    )
+    
+    test_dataset = TranslationDataset(
+        processed_data['test']['src'],
+        processed_data['test']['tgt']
+    )
+    
+    print(f"\nKích thước datasets:")
+    print(f"  - Train: {len(train_dataset)}")
+    print(f"  - Validation: {len(val_dataset)}")
+    print(f"  - Test: {len(test_dataset)}")
+    
+    # Tạo collate wrapper
+    collate_wrapper = CollateWrapper(pad_idx=0)
+    
+    # Tạo dataloaders
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        collate_fn=collate_wrapper,
+        num_workers=num_workers,
+        pin_memory=torch.cuda.is_available()  # Chỉ pin_memory nếu có GPU
+    )
+    
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        collate_fn=collate_wrapper,
+        num_workers=num_workers,
+        pin_memory=torch.cuda.is_available()
+    )
+    
+    test_loader = DataLoader(
+        test_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        collate_fn=collate_wrapper,
+        num_workers=num_workers,
+        pin_memory=torch.cuda.is_available()
+    )
+    
+    print(f"\nSố batch mỗi epoch:")
+    print(f"  - Train: {len(train_loader)}")
+    print(f"  - Validation: {len(val_loader)}")
+    print(f"  - Test: {len(test_loader)}")
+    
+    return train_loader, val_loader, test_loader
+
+# ============================================================================
+# 4. BUCKET SAMPLER - TỐI ƯU HƠN (OPTIONAL)
+# ============================================================================
+
+class BucketSampler(torch.utils.data.Sampler):
+    """
+    Sampler nhóm các câu có độ dài tương tự vào cùng batch
+    Giảm thiểu padding, tăng tốc huấn luyện
+    
+    TỐI ƯU: Chỉ sort một phần dataset để tránh lag với dataset lớn
+    """
+    def __init__(self, data_source, batch_size, sort_key=lambda x: len(x)):
+        self.data_source = data_source
+        self.batch_size = batch_size
+        self.sort_key = sort_key
+        
+    def __iter__(self):
+        # Sort toàn bộ dataset theo độ dài như trước
+        indices = list(range(len(self.data_source)))
+        lengths = [self.sort_key(self.data_source[i]) for i in indices]
+        all_indices = [i for i, _ in sorted(zip(indices, lengths), key=lambda x: x[1])]
+        
+        # Chia thành các batch
+        batches = [all_indices[i:i+self.batch_size] 
+                  for i in range(0, len(all_indices), self.batch_size)]
+        
+        # Shuffle thứ tự các batch (không shuffle trong batch)
+        np.random.shuffle(batches)
+        
+        # Flatten
+        for batch in batches:
+            for idx in batch:
+                yield idx
+    
+    def __len__(self):
+        return len(self.data_source)
+
+def create_dataloaders_with_bucketing(processed_data, batch_size=32, num_workers=0, val_batch_size=None):
+    """
+    Tạo DataLoader với BucketSampler để tối ưu padding
+    
+    IMPORTANT: num_workers phải = 0 trên Windows
+    
+    Args:
+        processed_data: Dictionary chứa train/validation/test data
+        batch_size: Batch size cho training (mặc định 32)
+        num_workers: Số workers cho DataLoader (mặc định 0)
+        val_batch_size: Batch size cho validation (mặc định = batch_size // 2 để tránh OOM)
+    """
+    # Validation batch size nhỏ hơn để tránh OOM
+    if val_batch_size is None:
+        val_batch_size = max(1, batch_size // 2)  # Giảm 50% so với train
+    print("\n" + "="*70)
+    print("TẠO DATALOADERS VỚI BUCKET SAMPLING")
+    print("="*70)
+    
+    # Tạo datasets
+    train_dataset = TranslationDataset(
+        processed_data['train']['src'],
+        processed_data['train']['tgt']
+    )
+    
+    val_dataset = TranslationDataset(
+        processed_data['validation']['src'],
+        processed_data['validation']['tgt']
+    )
+    
+    test_dataset = TranslationDataset(
+        processed_data['test']['src'],
+        processed_data['test']['tgt']
+    )
+    
+    # Tạo BucketSampler cho train - Sort toàn bộ dataset theo độ dài
+    train_sampler = BucketSampler(
+        train_dataset.src_data,
+        batch_size=batch_size,
+        sort_key=lambda x: len(x)
+    )
+    
+    # Tạo collate wrapper
+    collate_wrapper = CollateWrapper(pad_idx=0)
+    
+    # Tạo dataloaders với tối ưu hóa
+    use_gpu = torch.cuda.is_available()
+    pin_memory = use_gpu and num_workers > 0
+    persistent_workers = num_workers > 0
+    prefetch_factor = 2 if num_workers > 0 else None
+    
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        sampler=train_sampler,
+        collate_fn=collate_wrapper,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        persistent_workers=persistent_workers,
+        prefetch_factor=prefetch_factor
+    )
+    
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=val_batch_size,  # Dùng batch size nhỏ hơn cho validation
+        shuffle=False,
+        collate_fn=collate_wrapper,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        persistent_workers=persistent_workers,
+        prefetch_factor=prefetch_factor
+    )
+    
+    test_loader = DataLoader(
+        test_dataset,
+        batch_size=val_batch_size,  # Test cũng dùng batch size nhỏ hơn
+        shuffle=False,
+        collate_fn=collate_wrapper,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        persistent_workers=persistent_workers,
+        prefetch_factor=prefetch_factor
+    )
+    
+    print(f"✓ Đã tạo DataLoaders với Bucket Sampling")
+    print(f"  Train batch size: {batch_size}")
+    print(f"  Val/Test batch size: {val_batch_size} (giảm để tránh OOM)")
+    print(f"  Bucket Sampling giúp giảm padding, tăng tốc ~15-20%")
+    if num_workers == 0:
+        print(f"  ⚠️  num_workers=0 (Windows compatibility mode)")
+    
+    return train_loader, val_loader, test_loader
+
+# ============================================================================
+# 5. HELPER FUNCTIONS
+# ============================================================================
+
+def load_data_and_vocab(use_shared_vocab=True, use_bidirectional=False):
+    """
+    Load dữ liệu và vocabulary với Shared Vocabulary
+    
+    Args:
+        use_shared_vocab: Luôn True (giữ để tương thích)
+        use_bidirectional: Nếu True, load bidirectional dataset (có cả vi→en và en→vi)
+        
+    Returns:
+        processed_data: Dict với keys 'train', 'validation', 'test'
+        vi_vocab: Shared vocabulary wrapper (Vietnamese)
+        en_vocab: Shared vocabulary wrapper (English, cùng tokenizer với vi_vocab)
+    """
+    print("Đang load shared vocabulary data...")
+    
+    from shared_vocab_utils import (
+        load_shared_processed_data,
+        create_shared_vocab_wrapper
+    )
+    
+    processed_data = load_shared_processed_data(use_bidirectional=use_bidirectional)
+    
+    # Convert format từ list of tuples sang dict với 'src' và 'tgt'
+    # Format: {'train': [(src_ids, tgt_ids), ...]}
+    # → {'train': {'src': [src_ids, ...], 'tgt': [tgt_ids, ...]}}
+    converted_data = {}
+    for split in ['train', 'validation', 'test']:
+        if split in processed_data:
+            src_list = [item[0] for item in processed_data[split]]
+            tgt_list = [item[1] for item in processed_data[split]]
+            converted_data[split] = {
+                'src': src_list,
+                'tgt': tgt_list
+            }
+        else:
+            converted_data[split] = {'src': [], 'tgt': []}
+    
+    # Load shared vocab wrappers
+    vi_vocab, en_vocab = create_shared_vocab_wrapper()
+    
+    print("✓ Đã load shared vocabulary data!")
+    
+    return converted_data, vi_vocab, en_vocab
+
+def test_dataloader(loader, vi_vocab, en_vocab, num_batches=2):
+    """
+    Test dataloader và xem dữ liệu
+    """
+    print("\n" + "="*70)
+    print("KIỂM TRA DATALOADER")
+    print("="*70)
+    
+    for i, (src, tgt, src_len, tgt_len) in enumerate(loader):
+        if i >= num_batches:
+            break
+            
+        print(f"\nBatch {i+1}:")
+        print(f"  Source shape: {src.shape}")
+        print(f"  Target shape: {tgt.shape}")
+        print(f"  Source lengths: {src_len[:5].tolist()}...")
+        print(f"  Target lengths: {tgt_len[:5].tolist()}...")
+        
+        # Decode ví dụ đầu tiên
+        print(f"\n  Ví dụ đầu tiên trong batch:")
+        src_text = vi_vocab.decode(src[0].tolist())
+        tgt_text = en_vocab.decode(tgt[0].tolist())
+        print(f"    VI: {src_text}")
+        print(f"    EN: {tgt_text}")
+
+# ============================================================================
+# 6. MAIN
+# ============================================================================
+
+if __name__ == "__main__":
+    # Load dữ liệu
+    processed_data, vi_vocab, en_vocab = load_data_and_vocab()
+    
+    # Tạo dataloaders thông thường
+    print("\n" + "="*70)
+    print("PHƯƠNG ÁN 1: DATALOADER THÔNG THƯỜNG")
+    print("="*70)
+    train_loader, val_loader, test_loader = create_dataloaders(
+        processed_data, 
+        batch_size=32,
+        num_workers=0  # 0 cho Windows
+    )
+    
+    # Test
+    test_dataloader(train_loader, vi_vocab, en_vocab, num_batches=2)
+    
+    # Tạo dataloaders với bucketing (khuyên dùng)
+    print("\n" + "="*70)
+    print("PHƯƠNG ÁN 2: DATALOADER VỚI BUCKET SAMPLING (KHUYÊN DÙNG)")
+    print("="*70)
+    train_loader_bucket, val_loader_bucket, test_loader_bucket = create_dataloaders_with_bucketing(
+        processed_data,
+        batch_size=32,
+        num_workers=0  # 0 cho Windows
+    )
+    
+    # Test
+    test_dataloader(train_loader_bucket, vi_vocab, en_vocab, num_batches=2)
+    
+    print("\n" + "="*70)
+    print("✓ HOÀN TẤT TẠO DATALOADER!")
+    print("="*70)
+    print("\nGợi ý:")
+    print("  - Sử dụng batch_size=32 hoặc 64 tùy GPU")
+    print("  - Sử dụng BucketSampler để tối ưu tốc độ")
+    print("  - num_workers=0 trên Windows (multiprocessing issue)")
+    print("  - num_workers=2-4 trên Linux/Mac để tăng tốc")

src/encode_mtet_bidirectional.py

@@ -0,0 +1,240 @@
+"""
+Encode mtet_bidirectional.csv với shared vocabulary
+Chia train/validation/test và lưu thành processed data
+"""
+
+import pandas as pd
+import pickle
+import json
+from pathlib import Path
+from tqdm import tqdm
+from tokenizers import Tokenizer
+import sys
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+PROCESSED_DATA_DIR = PROJECT_ROOT / 'data' / 'processed'
+INPUT_FILE = PROJECT_ROOT / 'mtet_bidirectional.csv'
+
+def load_shared_tokenizer():
+    """Load shared tokenizer"""
+    tokenizer_path = PROCESSED_DATA_DIR / 'tokenizer_shared.json'
+    if not tokenizer_path.exists():
+        raise FileNotFoundError(
+            f"Không tìm thấy tokenizer_shared.json tại {tokenizer_path}!\n"
+            f"Vui lòng chạy: python src/1_build_shared_vocab.py trước"
+        )
+    return Tokenizer.from_file(str(tokenizer_path))
+
+def encode_csv_to_processed_data(
+    csv_file,
+    tokenizer,
+    train_ratio=0.95,
+    val_ratio=0.025,
+    test_ratio=0.025,
+    max_length=150,
+    chunk_size=100000
+):
+    """
+    Encode CSV file thành processed data và chia train/val/test
+    
+    Args:
+        csv_file: Path to CSV file (có cột src và tgt)
+        tokenizer: Shared tokenizer
+        train_ratio: Tỷ lệ train (mặc định 0.95)
+        val_ratio: Tỷ lệ validation (mặc định 0.025)
+        test_ratio: Tỷ lệ test (mặc định 0.025)
+        max_length: Độ dài tối đa của sequence (mặc định 150)
+        chunk_size: Kích thước chunk để xử lý
+        
+    Returns:
+        processed_data: Dict với keys 'train', 'validation', 'test'
+    """
+    print("="*70)
+    print("ENCODE MTET BIDIRECTIONAL DATASET")
+    print("="*70)
+    print(f"Input: {csv_file}")
+    print(f"Train: {train_ratio*100:.1f}% | Val: {val_ratio*100:.1f}% | Test: {test_ratio*100:.1f}%")
+    print(f"Max length: {max_length} tokens")
+    print("="*70 + "\n")
+    
+    # Kiểm tra file
+    if not csv_file.exists():
+        raise FileNotFoundError(f"Không tìm thấy file: {csv_file}")
+    
+    # Lấy special token IDs
+    SOS_ID = tokenizer.token_to_id("<sos>")
+    EOS_ID = tokenizer.token_to_id("<eos>")
+    EN2VI_ID = tokenizer.token_to_id("<en2vi>")
+    VI2EN_ID = tokenizer.token_to_id("<vi2en>")
+    
+    if SOS_ID is None or EOS_ID is None:
+        raise ValueError("Không tìm thấy <sos> hoặc <eos> trong tokenizer!")
+    if EN2VI_ID is None or VI2EN_ID is None:
+        raise ValueError("Không tìm thấy <en2vi> hoặc <vi2en> trong tokenizer!")
+    
+    # Đếm tổng số dòng
+    print("📊 Đang đếm số dòng...")
+    total_rows = sum(1 for _ in open(csv_file, 'r', encoding='utf-8')) - 1  # Trừ header
+    print(f"✓ Tổng số dòng: {total_rows:,}\n")
+    
+    # Encode tất cả data
+    print("🔄 Đang encode data...")
+    all_data = []
+    skipped = 0
+    
+    # Đọc và encode theo chunks
+    for chunk_df in tqdm(
+        pd.read_csv(csv_file, chunksize=chunk_size, encoding='utf-8'),
+        total=(total_rows // chunk_size) + 1,
+        desc="Encoding chunks"
+    ):
+        for _, row in chunk_df.iterrows():
+            src_text = str(row['src']).strip()
+            tgt_text = str(row['tgt']).strip()
+            
+            if not src_text or not tgt_text or src_text == 'nan' or tgt_text == 'nan':
+                skipped += 1
+                continue
+            
+            try:
+                # Detect direction: kiểm tra xem src là tiếng Anh hay tiếng Việt
+                # Phát hiện đơn giản: nếu có ký tự tiếng Việt (Unicode > 0x0100) thì là tiếng Việt
+                has_vietnamese_chars = any(ord(c) >= 0x0100 for c in src_text)
+                direction_id = VI2EN_ID if has_vietnamese_chars else EN2VI_ID
+                
+                # Encode
+                src_encoded = tokenizer.encode(src_text)
+                tgt_encoded = tokenizer.encode(tgt_text)
+                
+                src_ids = src_encoded.ids
+                tgt_ids = tgt_encoded.ids
+                
+                # Thêm SOS, direction token, và EOS
+                src_full = [SOS_ID, direction_id] + src_ids + [EOS_ID]
+                tgt_full = [SOS_ID] + tgt_ids + [EOS_ID]
+                
+                # Filter quá dài
+                if len(src_full) > max_length:
+                    src_full = src_full[:max_length-1] + [EOS_ID]
+                if len(tgt_full) > max_length:
+                    tgt_full = tgt_full[:max_length-1] + [EOS_ID]
+                
+                # Bỏ qua nếu quá ngắn (chỉ có SOS, direction token và EOS)
+                if len(src_full) <= 3 or len(tgt_full) <= 2:
+                    skipped += 1
+                    continue
+                
+                all_data.append((src_full, tgt_full))
+                
+            except Exception as e:
+                skipped += 1
+                continue
+    
+    print(f"\n✓ Đã encode {len(all_data):,} cặp câu")
+    if skipped > 0:
+        print(f"⚠️  Đã bỏ qua {skipped:,} cặp câu (lỗi hoặc quá ngắn)")
+    
+    # Chia train/val/test
+    print("\n📊 Đang chia train/validation/test...")
+    import random
+    random.seed(42)
+    random.shuffle(all_data)
+    
+    total = len(all_data)
+    train_end = int(total * train_ratio)
+    val_end = train_end + int(total * val_ratio)
+    
+    train_data = all_data[:train_end]
+    val_data = all_data[train_end:val_end]
+    test_data = all_data[val_end:]
+    
+    processed_data = {
+        'train': train_data,
+        'validation': val_data,
+        'test': test_data
+    }
+    
+    print(f"✓ Train: {len(train_data):,} cặp câu ({len(train_data)/total*100:.1f}%)")
+    print(f"✓ Validation: {len(val_data):,} cặp câu ({len(val_data)/total*100:.1f}%)")
+    print(f"✓ Test: {len(test_data):,} cặp câu ({len(test_data)/total*100:.1f}%)")
+    
+    return processed_data
+
+def main():
+    """Main function"""
+    import argparse
+    
+    parser = argparse.ArgumentParser(description='Encode mtet_bidirectional.csv')
+    parser.add_argument('--input', type=str, default=None,
+                       help='File CSV input (mặc định: mtet_bidirectional.csv)')
+    parser.add_argument('--output', type=str, default=None,
+                       help='File output (mặc định: processed_data_mtet_bidirectional.pkl)')
+    parser.add_argument('--train_ratio', type=float, default=0.95,
+                       help='Tỷ lệ train (mặc định: 0.95)')
+    parser.add_argument('--val_ratio', type=float, default=0.025,
+                       help='Tỷ lệ validation (mặc định: 0.025)')
+    parser.add_argument('--test_ratio', type=float, default=0.025,
+                       help='Tỷ lệ test (mặc định: 0.025)')
+    parser.add_argument('--max_length', type=int, default=150,
+                       help='Độ dài tối đa sequence (mặc định: 150)')
+    parser.add_argument('--chunk_size', type=int, default=100000,
+                       help='Kích thước chunk (mặc định: 100000)')
+    
+    args = parser.parse_args()
+    
+    # Paths
+    input_file = Path(args.input) if args.input else INPUT_FILE
+    output_file = args.output if args.output else (PROCESSED_DATA_DIR / 'processed_data_mtet_bidirectional.pkl')
+    output_file = Path(output_file)
+    
+    # Kiểm tra ratios
+    if abs(args.train_ratio + args.val_ratio + args.test_ratio - 1.0) > 0.01:
+        print("⚠️  Warning: Tổng các tỷ lệ không bằng 1.0, sẽ tự động điều chỉnh")
+        total = args.train_ratio + args.val_ratio + args.test_ratio
+        args.train_ratio /= total
+        args.val_ratio /= total
+        args.test_ratio /= total
+    
+    # Load tokenizer
+    print("📚 Đang load shared tokenizer...")
+    tokenizer = load_shared_tokenizer()
+    vocab_size = tokenizer.get_vocab_size()
+    print(f"✓ Vocab size: {vocab_size:,}\n")
+    
+    # Encode và chia data
+    processed_data = encode_csv_to_processed_data(
+        csv_file=input_file,
+        tokenizer=tokenizer,
+        train_ratio=args.train_ratio,
+        val_ratio=args.val_ratio,
+        test_ratio=args.test_ratio,
+        max_length=args.max_length,
+        chunk_size=args.chunk_size
+    )
+    
+    # Lưu file
+    print(f"\n💾 Đang lưu processed data...")
+    PROCESSED_DATA_DIR.mkdir(parents=True, exist_ok=True)
+    
+    with open(output_file, 'wb') as f:
+        pickle.dump(processed_data, f)
+    
+    print(f"✓ Đã lưu vào: {output_file}")
+    
+    # Thống kê
+    print("\n" + "="*70)
+    print("THỐNG KÊ")
+    print("="*70)
+    total = sum(len(v) for v in processed_data.values())
+    for split in ['train', 'validation', 'test']:
+        count = len(processed_data[split])
+        print(f"{split.capitalize():12s}: {count:>10,} cặp câu ({count/total*100:>5.1f}%)")
+    print(f"{'Tổng':12s}: {total:>10,} cặp câu")
+    print("="*70)
+    
+    print(f"\n✅ Hoàn thành! File đã sẵn sàng để train.")
+    print(f"   Sử dụng: python src/main_pipeline.py --use_mtet_bidirectional")
+
+if __name__ == '__main__':
+    main()
+

src/finetune.py

@@ -0,0 +1,520 @@
+"""
+FINETUNING SCRIPT
+Script để finetune model đã được train trước đó
+"""
+
+import torch
+import argparse
+import json
+from pathlib import Path
+import sys
+
+# Import modules
+from dataloader_module import load_data_and_vocab, create_dataloaders_with_bucketing
+from complete_transformer import create_model, print_model_info
+from training_module import train_model, plot_training_history, load_checkpoint
+from shared_vocab_utils import load_shared_vocab_info
+
+# ============================================================================
+# CONFIGURATION
+# ============================================================================
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DATA_DIR = PROJECT_ROOT / 'data'
+CHECKPOINT_DIR = PROJECT_ROOT / 'checkpoints'
+RESULTS_DIR = PROJECT_ROOT / 'results'
+PROCESSED_DATA_DIR = DATA_DIR / 'processed'
+
+class FinetuneConfig:
+    """Configuration cho finetuning"""
+    
+    def __init__(self):
+        # Paths
+        self.checkpoint_dir = CHECKPOINT_DIR
+        self.results_dir = RESULTS_DIR
+        
+        # Model
+        self.model_size = 'custom_25m'  # Phải khớp với checkpoint
+        self.use_weight_tying = True
+        
+        # Finetuning settings (Transfer Learning cổ điển)
+        self.checkpoint_path = None  # Sẽ được set từ args
+        self.learning_rate = 1e-5  # LR rất nhỏ (1e-5 hoặc 5e-6) - nhỏ hơn 10-100 lần training ban đầu
+        self.num_epochs = 5  # 5-10 epochs cho Transfer Learning
+        self.batch_size = 128
+        self.grad_accum_steps = 2
+        self.label_smoothing = 0.1
+        self.precision = 'bf16'  # 'bf16', 'amp', hoặc 'fp32'
+        self.dropout = None  # None = giữ nguyên, hoặc set giá trị mới (ví dụ: 0.2 nếu overfitting)
+        
+        # Freeze layers (tùy chọn)
+        self.freeze_encoder = False
+        self.freeze_decoder = False
+        self.freeze_embeddings = False
+        
+        # Device
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        
+        # Num workers
+        import platform
+        self.num_workers = 0 if platform.system() == 'Windows' else 2
+        
+        # Checkpoint
+        self.save_every = 1
+
+# ============================================================================
+# FREEZE LAYERS
+# ============================================================================
+
+def freeze_layers(model, freeze_encoder=False, freeze_decoder=False, freeze_embeddings=False):
+    """
+    Freeze một số layers của model
+    
+    Args:
+        model: Transformer model
+        freeze_encoder: Freeze toàn bộ encoder
+        freeze_decoder: Freeze toàn bộ decoder
+        freeze_embeddings: Freeze embedding layers
+    """
+    if freeze_embeddings:
+        print("→ Freezing embedding layers...")
+        if hasattr(model, 'shared_embedding'):
+            for param in model.shared_embedding.parameters():
+                param.requires_grad = False
+        print("  ✓ Embeddings frozen")
+    
+    if freeze_encoder:
+        print("→ Freezing encoder...")
+        for param in model.encoder.parameters():
+            param.requires_grad = False
+        print("  ✓ Encoder frozen")
+    
+    if freeze_decoder:
+        print("→ Freezing decoder...")
+        for param in model.decoder.parameters():
+            param.requires_grad = False
+        print("  ✓ Decoder frozen")
+    
+    # Đếm số parameters trainable
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    frozen_params = total_params - trainable_params
+    
+    print(f"\n  Total parameters: {total_params:,}")
+    print(f"  Trainable: {trainable_params:,} ({100*trainable_params/total_params:.1f}%)")
+    print(f"  Frozen: {frozen_params:,} ({100*frozen_params/total_params:.1f}%)")
+
+# ============================================================================
+# MAIN FINETUNING FUNCTION
+# ============================================================================
+
+def finetune_model(
+    checkpoint_path,
+    model_size='custom_25m',
+    learning_rate=1e-5,  # LR rất nhỏ cho Transfer Learning
+    num_epochs=5,
+    batch_size=128,
+    grad_accum_steps=2,
+    precision='bf16',
+    freeze_encoder=False,
+    freeze_decoder=False,
+    freeze_embeddings=False,
+    dropout=None,  # None = giữ nguyên, hoặc set giá trị mới
+    use_hospital_data=False,  # True = dùng dữ liệu Hospital, False = dùng dữ liệu gốc
+    save_every=1
+):
+    """
+    Finetune model từ checkpoint
+    
+    Args:
+        checkpoint_path: Đường dẫn đến checkpoint
+        model_size: Kích thước model (phải khớp với checkpoint)
+        finetune_lr_factor: Hệ số nhân learning rate (0.1 = 10% của lr ban đầu)
+        num_epochs: Số epochs để finetune
+        batch_size: Batch size
+        grad_accum_steps: Gradient accumulation steps
+        warmup_steps: Warmup steps (thường ít hơn training ban đầu)
+        precision: 'bf16', 'amp', hoặc 'fp32'
+        freeze_encoder: Freeze encoder layers
+        freeze_decoder: Freeze decoder layers
+        freeze_embeddings: Freeze embedding layers
+        save_every: Lưu checkpoint mỗi N epochs
+    """
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    print("="*70)
+    print("FINETUNING MODEL - TRANSFER LEARNING")
+    print("="*70)
+    print(f"\nCheckpoint: {checkpoint_path}")
+    print(f"Model size: {model_size}")
+    print(f"Learning rate: {learning_rate} (rất nhỏ - Transfer Learning)")
+    print(f"Epochs: {num_epochs} (5-10 epochs cho Transfer Learning)")
+    print(f"Device: {device}")
+    if dropout is not None:
+        print(f"Dropout: {dropout} (điều chỉnh để tránh overfitting)")
+    print()
+    
+    # Load vocabulary info
+    print("→ Loading vocabulary...")
+    vocab_info = load_shared_vocab_info()
+    vocab_size = vocab_info['vocab_size']
+    pad_idx = vocab_info['pad_id']
+    print(f"  ✓ Vocab size: {vocab_size}")
+    print(f"  ✓ PAD ID: {pad_idx}")
+    
+    # Create model
+    print("\n→ Creating model...")
+    model, model_config = create_model(
+        src_vocab_size=vocab_size,
+        tgt_vocab_size=vocab_size,
+        model_size=model_size,
+        pad_idx=pad_idx,
+        use_shared_vocab=True,
+        use_weight_tying=True
+    )
+    
+    # Điều chỉnh dropout nếu cần (để tránh overfitting)
+    if dropout is not None:
+        print(f"→ Điều chỉnh dropout: {dropout}")
+        # Điều chỉnh dropout trong encoder và decoder
+        for module in model.modules():
+            if isinstance(module, torch.nn.Dropout):
+                module.p = dropout
+        print(f"  ✓ Dropout đã được set thành {dropout}")
+    
+    model = model.to(device)
+    print("  ✓ Model created")
+    
+    # Load checkpoint
+    print(f"\n→ Loading checkpoint: {checkpoint_path}")
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    
+    # Print checkpoint info
+    checkpoint_epoch = checkpoint.get('epoch', 'N/A')
+    checkpoint_val_loss = checkpoint.get('val_loss', 'N/A')
+    print(f"  ✓ Checkpoint loaded")
+    print(f"    - Epoch: {checkpoint_epoch}")
+    print(f"    - Val Loss: {checkpoint_val_loss:.4f}" if isinstance(checkpoint_val_loss, float) else f"    - Val Loss: {checkpoint_val_loss}")
+    
+    # Freeze layers nếu cần
+    if freeze_encoder or freeze_decoder or freeze_embeddings:
+        print("\n→ Freezing layers...")
+        freeze_layers(model, freeze_encoder, freeze_decoder, freeze_embeddings)
+    
+    # Load data (có thể dùng dữ liệu Hospital hoặc dữ liệu gốc)
+    print("\n→ Loading data...")
+    
+    # Kiểm tra xem có dữ liệu Hospital không
+    hospital_data_path = PROCESSED_DATA_DIR / 'hospital_data_encoded.pkl'
+    use_hospital = use_hospital_data and hospital_data_path.exists()
+    
+    if use_hospital:
+        print("  → Sử dụng dữ liệu Hospital để finetune")
+        import pickle
+        with open(hospital_data_path, 'rb') as f:
+            processed_data = pickle.load(f)
+        
+        # Load vocab (dùng shared vocab)
+        from shared_vocab_utils import create_shared_vocab_wrapper
+        vi_vocab, en_vocab = create_shared_vocab_wrapper()
+        
+        print(f"  ✓ Hospital Train: {len(processed_data['train']['src'])} cặp câu")
+        print(f"  ✓ Hospital Validation: {len(processed_data['validation']['src'])} cặp câu")
+        if 'test' in processed_data:
+            print(f"  ✓ Hospital Test: {len(processed_data['test']['src'])} cặp câu")
+    else:
+        print("  → Sử dụng dữ liệu gốc")
+        processed_data, vi_vocab, en_vocab = load_data_and_vocab(use_shared_vocab=True)
+        print(f"  ✓ Train: {len(processed_data['train']['src'])} cặp câu")
+        print(f"  ✓ Validation: {len(processed_data['validation']['src'])} cặp câu")
+    
+    # Create dataloaders
+    print("\n→ Creating dataloaders...")
+    train_loader, val_loader, test_loader = create_dataloaders_with_bucketing(
+        processed_data,
+        batch_size=batch_size,
+        num_workers=0 if device.type == 'cpu' else 2
+    )
+    print("  ✓ Dataloaders created")
+    
+    # Setup optimizer với learning rate nhỏ hơn
+    # TransformerLRScheduler sẽ nhân với factor, nhưng chúng ta cần điều chỉnh base lr
+    # Thay vì thay đổi scheduler, chúng ta sẽ dùng finetune_lr_factor trong scheduler
+    
+    # Train với learning rate cố định rất nhỏ (Transfer Learning cổ điển)
+    print("\n" + "="*70)
+    print("BẮT ĐẦU FINETUNING - TRANSFER LEARNING")
+    print("="*70)
+    print(f"Phương pháp: Transfer Learning cổ điển (không dùng QLoRA/Unsloth)")
+    print(f"Learning rate: {learning_rate} (rất nhỏ - nhỏ hơn 10-100 lần training ban đầu)")
+    print(f"Optimizer: Adam mới tinh (không load từ checkpoint)")
+    print("="*70 + "\n")
+    
+    # QUAN TRỌNG: Optimizer mới tinh (không load từ checkpoint)
+    # Đây là điểm khác biệt với resume training
+    import torch.optim as optim
+    
+    # Chỉ optimize các parameters có requires_grad=True
+    trainable_params = [p for p in model.parameters() if p.requires_grad]
+    
+    # Setup optimizer với learning rate cố định rất nhỏ
+    optimizer = optim.Adam(
+        trainable_params,
+        lr=learning_rate,  # LR cố định rất nhỏ cho Transfer Learning
+        betas=(0.9, 0.98),
+        eps=1e-9
+    )
+    
+    # Scheduler đơn giản: StepLR hoặc CosineAnnealingLR
+    # Hoặc không dùng scheduler, chỉ dùng LR cố định
+    from torch.optim.lr_scheduler import StepLR
+    scheduler = StepLR(optimizer, step_size=num_epochs, gamma=0.5)  # Giảm 50% mỗi num_epochs
+    # Hoặc không dùng scheduler: scheduler = None và không gọi scheduler.step()
+    
+    # Setup loss
+    from training_module import LabelSmoothingLoss
+    criterion = LabelSmoothingLoss(
+        vocab_size=vocab_size,
+        pad_idx=pad_idx,
+        smoothing=0.1
+    )
+    
+    # Training history (bắt đầu mới cho finetune)
+    history = {
+        'train_loss': [],
+        'train_ppl': [],
+        'val_loss': [],
+        'val_ppl': [],
+        'lr': []
+    }
+    
+    # Lấy best_val_loss từ checkpoint gốc để so sánh
+    original_val_loss = checkpoint.get('val_loss', float('inf'))
+    best_val_loss = original_val_loss
+    print(f"→ Validation loss của checkpoint gốc: {original_val_loss:.4f}")
+    
+    # Import training functions
+    from training_module import train_epoch, validate, calculate_perplexity
+    import time
+    from tqdm import tqdm
+    
+    use_amp = precision == 'amp' and device.type == 'cuda'
+    use_bf16 = precision == 'bf16' and device.type == 'cuda'
+    from torch.cuda.amp import GradScaler
+    scaler = GradScaler(enabled=use_amp and not use_bf16) if use_amp and not use_bf16 else None
+    
+    print("="*70)
+    print("FINETUNING - TRANSFER LEARNING")
+    print("="*70)
+    print(f"Device: {device}")
+    print(f"Number of epochs: {num_epochs}")
+    print(f"Learning rate: {learning_rate} (cố định)")
+    print(f"Grad accumulation steps: {grad_accum_steps}")
+    print(f"Original val loss: {original_val_loss:.4f}")
+    if use_bf16:
+        print(f"Precision: BF16")
+    elif use_amp:
+        print(f"Precision: AMP")
+    else:
+        print(f"Precision: FP32")
+    print("="*70 + "\n")
+    
+    for epoch in range(1, num_epochs + 1):
+        epoch_start_time = time.time()
+        
+        # Training
+        train_loss, train_ppl = train_epoch(
+            model,
+            train_loader,
+            optimizer,
+            scheduler,
+            criterion,
+            device,
+            epoch,
+            grad_accum_steps=grad_accum_steps,
+            use_amp=use_amp,
+            use_bf16=use_bf16,
+            scaler=scaler,
+        )
+        
+        # Validation
+        val_loss, val_ppl = validate(
+            model,
+            val_loader,
+            criterion,
+            device,
+            use_amp=use_amp,
+            use_bf16=use_bf16,
+        )
+        
+        epoch_time = time.time() - epoch_start_time
+        
+        # Save history
+        history['train_loss'].append(train_loss)
+        history['train_ppl'].append(train_ppl)
+        history['val_loss'].append(val_loss)
+        history['val_ppl'].append(val_ppl)
+        current_lr = scheduler.get_last_lr()[0] if scheduler else learning_rate
+        history['lr'].append(current_lr)
+        
+        # Save checkpoint
+        checkpoint_msg = ""
+        best_msg = ""
+        
+        if epoch % save_every == 0:
+            checkpoint_path_save = CHECKPOINT_DIR / f'finetune_checkpoint_epoch_{epoch}.pt'
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'scheduler_state_dict': scheduler.state_dict() if scheduler else None,
+                'train_loss': train_loss,
+                'val_loss': val_loss,
+                'history': history,
+                'finetune_config': {
+                    'checkpoint_path': str(checkpoint_path),
+                    'learning_rate': learning_rate,
+                    'model_size': model_size,
+                    'original_val_loss': original_val_loss,
+                    'method': 'transfer_learning_classic'
+                }
+            }, checkpoint_path_save)
+            checkpoint_msg = " | ✓ Saved checkpoint"
+        
+        # Save best model
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            best_model_path = CHECKPOINT_DIR / 'best_model_finetuned.pt'
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'val_loss': val_loss,
+                'val_ppl': val_ppl,
+                'finetune_config': {
+                    'checkpoint_path': str(checkpoint_path),
+                    'learning_rate': learning_rate,
+                    'model_size': model_size,
+                    'original_val_loss': original_val_loss,
+                    'improvement': original_val_loss - val_loss,
+                    'method': 'transfer_learning_classic'
+                }
+            }, best_model_path)
+            improvement = original_val_loss - val_loss
+            best_msg = f" | ✓ New best! (Cải thiện: {improvement:.4f})"
+        
+        current_lr = scheduler.get_last_lr()[0] if scheduler else learning_rate
+        print(f"Epoch {epoch}/{num_epochs} | Time: {epoch_time:.2f}s | Train Loss: {train_loss:.4f} | Train PPL: {train_ppl:.2f} | Val Loss: {val_loss:.4f} | Val PPL: {val_ppl:.2f} | LR: {current_lr:.6f}{checkpoint_msg}{best_msg}")
+        
+        # Step scheduler (nếu có)
+        if scheduler:
+            scheduler.step()
+    
+    print("\n" + "="*70)
+    print("HOÀN TẤT FINETUNING - TRANSFER LEARNING")
+    print("="*70)
+    print(f"Original validation loss: {original_val_loss:.4f}")
+    print(f"Best validation loss: {best_val_loss:.4f}")
+    improvement = original_val_loss - best_val_loss
+    if improvement > 0:
+        print(f"✓ Cải thiện: {improvement:.4f} ({improvement/original_val_loss*100:.2f}%)")
+    else:
+        print(f"⚠️  Không cải thiện (có thể cần điều chỉnh LR hoặc epochs)")
+    print(f"Best validation perplexity: {calculate_perplexity(best_val_loss):.2f}")
+    
+    # Plot history
+    finetune_plot_path = RESULTS_DIR / 'finetune_history.png'
+    plot_training_history(history, save_path=finetune_plot_path)
+    
+    # Save history
+    finetune_history_path = RESULTS_DIR / 'finetune_history.json'
+    with open(finetune_history_path, 'w') as f:
+        json.dump(history, f, indent=2)
+    
+    print(f"\n✓ Saved finetune history to {finetune_history_path}")
+    print(f"✓ Saved finetune plot to {finetune_plot_path}")
+    
+    return history
+
+# ============================================================================
+# MAIN
+# ============================================================================
+
+def main():
+    parser = argparse.ArgumentParser(description='Finetune Transformer Model')
+    
+    # Checkpoint
+    parser.add_argument('--checkpoint', type=str, 
+                       default='checkpoints/best_model.pt',
+                       help='Path to checkpoint to finetune from')
+    
+    # Model
+    parser.add_argument('--model_size', type=str, default='custom_25m',
+                       help='Model size (must match checkpoint)')
+    
+    # Finetuning settings (Transfer Learning cổ điển)
+    parser.add_argument('--learning_rate', type=float, default=1e-5,
+                       help='Learning rate (rất nhỏ: 1e-5 hoặc 5e-6, nhỏ hơn 10-100 lần training ban đầu)')
+    parser.add_argument('--epochs', type=int, default=5,
+                       help='Number of epochs (5-10 epochs cho Transfer Learning)')
+    parser.add_argument('--batch_size', type=int, default=128,
+                       help='Batch size')
+    parser.add_argument('--grad_accum_steps', type=int, default=2,
+                       help='Gradient accumulation steps')
+    parser.add_argument('--precision', type=str, default='bf16',
+                       choices=['bf16', 'amp', 'fp32'],
+                       help='Training precision')
+    parser.add_argument('--dropout', type=float, default=None,
+                       help='Dropout rate (None = giữ nguyên, hoặc set 0.2 nếu overfitting)')
+    parser.add_argument('--use_hospital', action='store_true',
+                       help='Use Hospital domain data for finetuning (requires prepare_hospital_data.py first)')
+    
+    # Freeze layers
+    parser.add_argument('--freeze_encoder', action='store_true',
+                       help='Freeze encoder layers (chỉ train decoder)')
+    parser.add_argument('--freeze_decoder', action='store_true',
+                       help='Freeze decoder layers (chỉ train encoder)')
+    parser.add_argument('--freeze_embeddings', action='store_true',
+                       help='Freeze embedding layers')
+    
+    args = parser.parse_args()
+    
+    # Resolve checkpoint path
+    checkpoint_path = Path(args.checkpoint)
+    if not checkpoint_path.is_absolute():
+        checkpoint_path = PROJECT_ROOT / checkpoint_path
+    
+    if not checkpoint_path.exists():
+        print(f"❌ Error: Checkpoint not found: {checkpoint_path}")
+        print(f"   Available checkpoints:")
+        for ckpt in CHECKPOINT_DIR.glob('*.pt'):
+            print(f"     - {ckpt}")
+        sys.exit(1)
+    
+    # Run finetuning
+    finetune_model(
+        checkpoint_path=str(checkpoint_path),
+        model_size=args.model_size,
+        learning_rate=args.learning_rate,
+        num_epochs=args.epochs,
+        batch_size=args.batch_size,
+        grad_accum_steps=args.grad_accum_steps,
+        precision=args.precision,
+        freeze_encoder=args.freeze_encoder,
+        freeze_decoder=args.freeze_decoder,
+        freeze_embeddings=args.freeze_embeddings,
+        dropout=args.dropout,
+        use_hospital_data=args.use_hospital,
+        save_every=1
+    )
+    
+    print("\n" + "="*70)
+    print("🎉 FINETUNING HOÀN TẤT!")
+    print("="*70)
+    print(f"\nBest finetuned model: {CHECKPOINT_DIR / 'best_model_finetuned.pt'}")
+    print(f"Training history: {RESULTS_DIR / 'finetune_history.json'}")
+    print(f"Training plot: {RESULTS_DIR / 'finetune_history.png'}")
+
+if __name__ == "__main__":
+    main()
+

src/finetune_hospital.py

@@ -0,0 +1,343 @@
+"""
+Finetune model với Hospital dataset (1M Hospital + 400k old data)
+"""
+
+import os
+# Set CUDA memory allocation để tránh fragmentation
+os.environ['PYTORCH_ALLOC_CONF'] = 'expandable_segments:True'
+
+import torch
+from pathlib import Path
+import sys
+import argparse
+
+# Add parent directory to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from src.complete_transformer import TransformerShared, get_model_config
+from src.shared_vocab_utils import load_shared_tokenizer, load_shared_vocab_info
+from src.dataloader_module import create_dataloaders_with_bucketing
+from src.training_module import train_model, load_checkpoint
+from src.gpu_safety import get_gpu_memory_info, check_and_adjust_batch_size, clear_gpu_cache
+import pickle
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+PROCESSED_DATA_DIR = PROJECT_ROOT / 'data' / 'processed'
+CHECKPOINT_DIR = PROJECT_ROOT / 'checkpoints'
+
+def load_hospital_finetune_data():
+    """Load processed hospital finetune data (đã làm sạch)"""
+    data_path = PROCESSED_DATA_DIR / "processed_data_hospital_finetune_cleaned.pkl"
+    
+    if not data_path.exists():
+        raise FileNotFoundError(
+            f"Không tìm thấy processed data tại {data_path}!\n"
+            f"Vui lòng chạy: python src/encode_mtet_bidirectional.py --input hospital_finetune_bidirectional_cleaned.csv --output data/processed/processed_data_hospital_finetune_cleaned.pkl trước"
+        )
+    
+    print(f"📂 Loading data từ: {data_path}")
+    with open(data_path, 'rb') as f:
+        data = pickle.load(f)
+    
+    # Convert format từ list of tuples sang dict với 'src' và 'tgt'
+    converted_data = {}
+    for split in ['train', 'validation', 'test']:
+        if split in data:
+            src_list = [item[0] for item in data[split]]
+            tgt_list = [item[1] for item in data[split]]
+            converted_data[split] = {
+                'src': src_list,
+                'tgt': tgt_list
+            }
+        else:
+            converted_data[split] = {'src': [], 'tgt': []}
+    
+    return converted_data
+
+def main():
+    parser = argparse.ArgumentParser(description="Finetune với Hospital dataset")
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        default="checkpoints/best_model.pt",
+        help="Checkpoint để load (mặc định: checkpoints/best_model.pt)",
+    )
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        default=5,
+        help="Số epochs (mặc định: 5)",
+    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=64,
+        help="Batch size (mặc định: 64)",
+    )
+    parser.add_argument(
+        "--model_size",
+        type=str,
+        default="custom_25m",
+        choices=["custom_25m", "base", "small"],
+        help="Kích thước model (mặc định: custom_25m)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda",
+        help="Device (cuda/cpu, mặc định: cuda)",
+    )
+    parser.add_argument(
+        "--precision",
+        type=str,
+        default="bf16",
+        choices=["fp32", "amp", "bf16"],
+        help="Precision (mặc định: bf16)",
+    )
+    parser.add_argument(
+        "--grad_accum_steps",
+        type=int,
+        default=2,
+        help="Gradient accumulation steps (mặc định: 2)",
+    )
+    parser.add_argument(
+        "--warmup_steps",
+        type=int,
+        default=2000,
+        help="Warmup steps (mặc định: 2000 steps cho finetune)",
+    )
+    parser.add_argument(
+        "--target_max_lr",
+        type=float,
+        default=1e-5,
+        help="Đỉnh learning rate cho finetune (mặc định: 1e-5)",
+    )
+    parser.add_argument(
+        "--label_smoothing",
+        type=float,
+        default=0.1,
+        help="Label smoothing (mặc định: 0.1)",
+    )
+    parser.add_argument(
+        "--checkpoint_dir",
+        type=str,
+        default=None,
+        help="Checkpoint directory (mặc định: checkpoints)",
+    )
+    parser.add_argument(
+        "--no_compile",
+        action="store_true",
+        help="Tắt torch.compile (mặc định: BẬT nếu grad_accum_steps=1)",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=None,
+        help="Số workers cho DataLoader (mặc định: 2 cho CUDA, 0 cho CPU)",
+    )
+    
+    args = parser.parse_args()
+    
+    # Device
+    device = torch.device(args.device)
+    if device.type == 'cuda' and not torch.cuda.is_available():
+        print("⚠️  CUDA không khả dụng, dùng CPU")
+        device = torch.device('cpu')
+    
+    print("="*70)
+    print("FINETUNE MODEL VỚI HOSPITAL DATASET")
+    print("="*70)
+    print(f"Device: {device}")
+    print(f"Model size: {args.model_size}")
+    print(f"Checkpoint: {args.checkpoint}")
+    
+    # Kiểm tra checkpoint
+    checkpoint_path = Path(args.checkpoint)
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(
+            f"Không tìm thấy checkpoint: {checkpoint_path}\n"
+            f"Vui lòng chỉ định checkpoint hợp lệ"
+        )
+    
+    # Kiểm tra GPU memory và điều chỉnh batch size nếu cần
+    if device.type == 'cuda':
+        clear_gpu_cache()
+        mem_info = get_gpu_memory_info()
+        if mem_info:
+            print(f"\n📊 GPU Memory Info:")
+            print(f"   Total: {mem_info['total']:.2f} GB")
+            print(f"   Free: {mem_info['free']:.2f} GB")
+            print(f"   Usage: {mem_info['usage_percent']:.1f}%")
+        
+        # Kiểm tra và điều chỉnh batch size
+        safe_batch_size, warning = check_and_adjust_batch_size(args.batch_size)
+        if warning:
+            print(f"\n{warning}")
+            response = input(f"\n❓ Có muốn giảm batch_size xuống {safe_batch_size} không? (y/n, mặc định: n): ").strip().lower()
+            if response == 'y':
+                args.batch_size = safe_batch_size
+                print(f"✓ Đã giảm batch_size xuống {safe_batch_size}")
+            else:
+                print(f"⚠️  Giữ nguyên batch_size={args.batch_size}, có thể gặp OOM")
+        else:
+            print(f"✓ Batch size {args.batch_size} an toàn với GPU memory hiện tại")
+    
+    print(f"\nBatch size: {args.batch_size}")
+    print(f"Gradient accumulation: {args.grad_accum_steps}")
+    effective_batch = args.batch_size * args.grad_accum_steps
+    print(f"Effective batch size: {effective_batch} ({args.batch_size} × {args.grad_accum_steps})")
+    print(f"Epochs: {args.num_epochs}")
+    print(f"Precision: {args.precision}")
+    print(f"Learning rate: {args.target_max_lr} (finetune - thấp hơn train từ đầu)")
+    print(f"Warmup steps: {args.warmup_steps}")
+    print()
+    print("🔄 FINETUNE SETUP:")
+    print("   - Dataset: 994k Hospital (đã làm sạch) + 400k old data (bidirectional)")
+    print("   - Dữ liệu đã được làm sạch: loại bỏ rác, lặp từ, đảm bảo alignment")
+    print("   - Learning rate thấp để tránh quên kiến thức cũ")
+    print("   - Warmup để ổn định training")
+    print("   - Label smoothing để tránh overfitting")
+    print("="*70 + "\n")
+    
+    # Load vocab info
+    print("📚 Loading vocabulary...")
+    vocab_info = load_shared_vocab_info()
+    vocab_size = vocab_info['vocab_size']
+    pad_idx = vocab_info['pad_id']
+    print(f"✓ Vocab size: {vocab_size:,}")
+    print(f"✓ Pad ID: {pad_idx}\n")
+    
+    # Load data
+    print("📂 Loading Hospital finetune data...")
+    processed_data = load_hospital_finetune_data()
+    
+    print(f"✓ Train: {len(processed_data['train']['src']):,} cặp câu")
+    print(f"✓ Validation: {len(processed_data['validation']['src']):,} cặp câu")
+    print(f"✓ Test: {len(processed_data['test']['src']):,} cặp câu\n")
+    
+    # Create dataloaders
+    print("🔄 Creating dataloaders...")
+    if args.num_workers is None:
+        num_workers = 2 if device.type == 'cuda' else 0
+    else:
+        num_workers = args.num_workers
+    train_loader, val_loader, test_loader = create_dataloaders_with_bucketing(
+        processed_data,
+        batch_size=args.batch_size,
+        num_workers=num_workers,
+    )
+    print(f"✓ Dataloaders created với BucketSampler (num_workers={num_workers})\n")
+    
+    # Create model
+    print("🔨 Creating model...")
+    model_config = get_model_config(args.model_size)
+    
+    model = TransformerShared(
+        vocab_size=vocab_size,
+        d_model=model_config['d_model'],
+        n_layers=model_config['n_layers'],
+        n_heads=model_config['n_heads'],
+        d_ff=model_config['d_ff'],
+        dropout=model_config['dropout'],
+        pad_idx=pad_idx,
+        use_weight_tying=True
+    )
+    
+    model = model.to(device)
+    
+    # Load checkpoint
+    print(f"📥 Loading checkpoint từ: {checkpoint_path}")
+    try:
+        model, checkpoint = load_checkpoint(model, str(checkpoint_path), device)
+        print(f"✓ Đã load checkpoint từ epoch {checkpoint.get('epoch', 'unknown')}")
+        print(f"✓ Best validation loss: {checkpoint.get('best_val_loss', checkpoint.get('val_loss', 'unknown'))}")
+    except Exception as e:
+        print(f"⚠️  Không thể load checkpoint: {e}")
+        print("   Sẽ train từ đầu")
+    
+    # Tối ưu: Compile model với torch.compile (PyTorch 2.0+)
+    compile_model = not args.no_compile and args.grad_accum_steps == 1
+    if compile_model and hasattr(torch, 'compile') and device.type == 'cuda':
+        print("🔧 Compiling model với torch.compile để tăng tốc...")
+        try:
+            model = torch.compile(model, mode='default')
+            print("✓ Model đã được compile\n")
+        except Exception as e:
+            print(f"⚠️  Không thể compile model: {e}\n")
+    else:
+        if args.grad_accum_steps > 1:
+            print("⚠️  torch.compile đã tắt (không tương thích với gradient accumulation)\n")
+        elif args.no_compile:
+            print("⚠️  torch.compile đã tắt (--no_compile)\n")
+    
+    # Count parameters
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    model_size_mb = total_params * 4 / (1024**2)
+    print(f"✓ Model created")
+    print(f"  - Total parameters: {total_params:,}")
+    print(f"  - Trainable parameters: {trainable_params:,}")
+    print(f"  - Model size: ~{model_size_mb:.2f} MB (float32)")
+    
+    # Kiểm tra lại GPU memory sau khi load model
+    if device.type == 'cuda':
+        clear_gpu_cache()
+        mem_info = get_gpu_memory_info()
+        if mem_info:
+            print(f"\n📊 GPU Memory sau khi load model:")
+            print(f"   Allocated: {mem_info['allocated']:.2f} GB")
+            print(f"   Reserved: {mem_info['reserved']:.2f} GB")
+            print(f"   Free: {mem_info['free']:.2f} GB")
+            print(f"   Usage: {mem_info['usage_percent']:.1f}%")
+    
+    print()
+    
+    # Checkpoint directory - Lưu vào thư mục riêng để không ghi đè checkpoint cũ
+    if args.checkpoint_dir:
+        checkpoint_dir = Path(args.checkpoint_dir)
+    else:
+        # Tự động tạo thư mục riêng cho finetune
+        checkpoint_dir = CHECKPOINT_DIR / 'finetune_hospital'
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    print(f"📁 Checkpoint sẽ được lưu tại: {checkpoint_dir}")
+    print(f"   (Không ghi đè checkpoint cũ trong {CHECKPOINT_DIR})\n")
+    
+    # Clear cache trước khi training
+    if device.type == 'cuda':
+        clear_gpu_cache()
+        print("✓ Đã clear GPU cache trước khi training\n")
+    
+    # Train
+    print("🚀 Starting finetune...")
+    print("="*70 + "\n")
+    
+    history = train_model(
+        model=model,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        num_epochs=args.num_epochs,
+        device=device,
+        d_model=model_config["d_model"],
+        warmup_steps=args.warmup_steps,
+        label_smoothing=args.label_smoothing,
+        grad_accum_steps=args.grad_accum_steps,
+        precision=args.precision,
+        checkpoint_dir=checkpoint_dir,
+        save_every=1,
+        resume_from=None,  # Không resume khi finetune
+        target_max_lr=args.target_max_lr,
+    )
+    
+    print("\n" + "="*70)
+    print("✅ FINETUNE HOÀN TẤT!")
+    print("="*70)
+    print(f"Checkpoints được lưu tại: {checkpoint_dir}")
+    print(f"Best model: {checkpoint_dir / 'best_model.pt'}")
+    print(f"\n💡 Lưu ý: Checkpoint cũ vẫn an toàn tại {CHECKPOINT_DIR}")
+    print(f"   - Model gốc: {CHECKPOINT_DIR / 'best_model.pt'}")
+    print(f"   - Model finetune: {checkpoint_dir / 'best_model.pt'}")
+    print("="*70)
+
+if __name__ == '__main__':
+    main()
+

src/gpu_safety.py

@@ -0,0 +1,197 @@
+"""
+Cơ chế an toàn GPU: Tự động điều chỉnh batch size để tránh OOM
+"""
+
+import torch
+import gc
+
+
+def get_gpu_memory_info():
+    """Lấy thông tin GPU memory"""
+    if not torch.cuda.is_available():
+        return None
+    
+    memory_allocated = torch.cuda.memory_allocated() / (1024**3)  # GB
+    memory_reserved = torch.cuda.memory_reserved() / (1024**3)  # GB
+    memory_total = torch.cuda.get_device_properties(0).total_memory / (1024**3)  # GB
+    memory_free = memory_total - memory_reserved
+    
+    return {
+        'allocated': memory_allocated,
+        'reserved': memory_reserved,
+        'total': memory_total,
+        'free': memory_free,
+        'usage_percent': (memory_reserved / memory_total) * 100
+    }
+
+
+def calculate_safe_batch_size(current_batch_size, model_size_mb, vocab_size, 
+                              max_seq_len=150, safety_margin=0.15):
+    """
+    Tính toán batch size an toàn dựa trên GPU memory còn lại
+    
+    Args:
+        current_batch_size: Batch size hiện tại
+        model_size_mb: Kích thước model (MB)
+        vocab_size: Kích thước vocabulary
+        max_seq_len: Độ dài sequence tối đa
+        safety_margin: Margin an toàn (15% mặc định)
+    
+    Returns:
+        safe_batch_size: Batch size an toàn
+    """
+    if not torch.cuda.is_available():
+        return current_batch_size
+    
+    mem_info = get_gpu_memory_info()
+    if mem_info is None:
+        return current_batch_size
+    
+    # Ước tính memory cần cho một batch
+    # Memory cho model: model_size_mb
+    # Memory cho activations: batch_size * seq_len * hidden_size * 4 bytes (float32)
+    # Memory cho gradients: tương tự activations
+    # Memory cho optimizer states: ~2x model size (Adam)
+    
+    # Ước tính đơn giản: mỗi batch cần ~50-100MB tùy vào seq_len
+    estimated_memory_per_batch_mb = 80  # Conservative estimate
+    
+    # Memory còn lại có thể dùng (trừ safety margin)
+    available_memory_gb = mem_info['free'] * (1 - safety_margin)
+    available_memory_mb = available_memory_gb * 1024
+    
+    # Trừ memory cho model và optimizer
+    model_memory_mb = model_size_mb * 3  # Model + gradients + optimizer states
+    usable_memory_mb = available_memory_mb - model_memory_mb
+    
+    if usable_memory_mb < 0:
+        # Không đủ memory, giảm batch size xuống tối thiểu
+        return max(1, current_batch_size // 4)
+    
+    # Tính batch size an toàn
+    safe_batch_size = int(usable_memory_mb / estimated_memory_per_batch_mb)
+    
+    # Giới hạn trong khoảng hợp lý
+    safe_batch_size = max(1, min(safe_batch_size, current_batch_size * 2))
+    
+    return safe_batch_size
+
+
+def check_and_adjust_batch_size(current_batch_size, threshold=0.85):
+    """
+    Kiểm tra GPU memory và đề xuất batch size mới nếu cần
+    
+    Args:
+        current_batch_size: Batch size hiện tại
+        threshold: Ngưỡng cảnh báo (85% mặc định)
+    
+    Returns:
+        new_batch_size: Batch size đề xuất (có thể giống hoặc nhỏ hơn)
+        warning: Cảnh báo nếu cần
+    """
+    if not torch.cuda.is_available():
+        return current_batch_size, None
+    
+    mem_info = get_gpu_memory_info()
+    if mem_info is None:
+        return current_batch_size, None
+    
+    usage_percent = mem_info['usage_percent']
+    warning = None
+    
+    if usage_percent > threshold * 100:
+        # GPU memory gần hết, đề xuất giảm batch size
+        reduction_factor = 0.75  # Giảm 25%
+        new_batch_size = max(1, int(current_batch_size * reduction_factor))
+        warning = (
+            f"⚠️  GPU memory usage: {usage_percent:.1f}% "
+            f"(Free: {mem_info['free']:.2f}GB / Total: {mem_info['total']:.2f}GB)\n"
+            f"   Đề xuất giảm batch_size từ {current_batch_size} xuống {new_batch_size}"
+        )
+        return new_batch_size, warning
+    elif usage_percent > 0.7 * 100:
+        # Cảnh báo sớm
+        warning = (
+            f"⚠️  GPU memory usage: {usage_percent:.1f}% "
+            f"(Free: {mem_info['free']:.2f}GB / Total: {mem_info['total']:.2f}GB)\n"
+            f"   Nên theo dõi, có thể cần giảm batch_size nếu tiếp tục tăng"
+        )
+        return current_batch_size, warning
+    
+    return current_batch_size, None
+
+
+def clear_gpu_cache():
+    """Xóa GPU cache để giải phóng memory"""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        gc.collect()
+
+
+def monitor_gpu_memory(interval=100):
+    """
+    Decorator để monitor GPU memory trong training loop
+    
+    Args:
+        interval: Kiểm tra mỗi N batches
+    """
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            # Chạy function và monitor
+            result = func(*args, **kwargs)
+            
+            # Kiểm tra memory sau mỗi interval batches
+            if hasattr(wrapper, 'batch_count'):
+                wrapper.batch_count += 1
+            else:
+                wrapper.batch_count = 1
+            
+            if wrapper.batch_count % interval == 0:
+                mem_info = get_gpu_memory_info()
+                if mem_info:
+                    print(f"[GPU Monitor] Memory: {mem_info['usage_percent']:.1f}% "
+                          f"({mem_info['free']:.2f}GB free)")
+            
+            return result
+        return wrapper
+    return decorator
+
+
+def handle_oom_error(current_batch_size, reduction_factor=0.75, min_batch_size=8):
+    """
+    Xử lý OOM error bằng cách giảm batch size
+    
+    Args:
+        current_batch_size: Batch size hiện tại
+        reduction_factor: Hệ số giảm (mặc định: 0.75 = giảm 25%)
+        min_batch_size: Batch size tối thiểu
+    
+    Returns:
+        new_batch_size: Batch size mới sau khi giảm
+    """
+    new_batch_size = max(min_batch_size, int(current_batch_size * reduction_factor))
+    
+    # Xóa cache
+    clear_gpu_cache()
+    
+    return new_batch_size
+
+
+def check_memory_spike(threshold=0.9):
+    """
+    Kiểm tra xem GPU memory có đang spike không
+    
+    Args:
+        threshold: Ngưỡng cảnh báo (90% mặc định)
+    
+    Returns:
+        is_spike: True nếu memory > threshold
+        mem_info: Thông tin memory
+    """
+    mem_info = get_gpu_memory_info()
+    if mem_info is None:
+        return False, None
+    
+    is_spike = mem_info['usage_percent'] > threshold * 100
+    return is_spike, mem_info
+

src/inference_evaluation.py

@@ -0,0 +1,563 @@
+"""
+INFERENCE & EVALUATION
+Greedy Search, Beam Search, BLEU Score evaluation
+"""
+
+import torch
+import torch.nn.functional as F
+from collections import Counter
+import math
+from tqdm import tqdm
+
+# ============================================================================
+# 1. GREEDY SEARCH DECODING
+# ============================================================================
+
+def greedy_decode(model, src, src_vocab, tgt_vocab, device, max_len=100, repetition_penalty=1.5, no_repeat_ngram_size=3):
+    """
+    Greedy Decoding - Chọn token có xác suất cao nhất mỗi bước với Repetition Penalty
+    
+    Args:
+        model: Transformer model
+        src: Source sequence [1, src_len] hoặc list of tokens
+        src_vocab: Source vocabulary
+        tgt_vocab: Target vocabulary
+        device: Device
+        max_len: Maximum length để generate
+        repetition_penalty: Penalty factor cho tokens lặp lại (>1.0 để giảm repetition)
+        no_repeat_ngram_size: Kích thước n-gram để tránh lặp lại (0 = tắt)
+        
+    Returns:
+        decoded_tokens: List of decoded token indices
+        decoded_sentence: Decoded sentence (string)
+    """
+    model.eval()
+    
+    # Nếu src là list, convert sang tensor
+    if isinstance(src, list):
+        src = torch.LongTensor([src]).to(device)
+    elif src.dim() == 1:
+        src = src.unsqueeze(0)
+    
+    with torch.no_grad():
+        # Encode source
+        encoder_output, src_mask = model.encode(src)
+        
+        # Khởi tạo target với <sos> token
+        tgt_tokens = [tgt_vocab.SOS_IDX]
+        
+        for step in range(max_len):
+            # Tạo target tensor
+            tgt = torch.LongTensor([tgt_tokens]).to(device)
+            
+            # Decode
+            output = model.decode(tgt, encoder_output, src_mask)
+            
+            # Lấy prediction cho token cuối cùng
+            next_token_logits = output[0, -1, :]
+            
+            # Áp dụng repetition penalty cho từng token
+            token_counts = {}
+            for token_id in tgt_tokens:
+                if token_id not in [tgt_vocab.SOS_IDX, tgt_vocab.EOS_IDX, tgt_vocab.PAD_IDX]:
+                    token_counts[token_id] = token_counts.get(token_id, 0) + 1
+            
+            # Áp dụng penalty: giảm logit của các tokens đã xuất hiện
+            for token_id, count in token_counts.items():
+                if count > 0 and token_id < len(next_token_logits):
+                    # Penalty tăng theo số lần lặp lại
+                    penalty = repetition_penalty ** count
+                    next_token_logits[token_id] = next_token_logits[token_id] / penalty
+            
+            # N-gram repetition penalty: tránh lặp cụm từ
+            if no_repeat_ngram_size > 0 and len(tgt_tokens) >= no_repeat_ngram_size:
+                # Lấy n-gram cuối cùng
+                last_ngram = tuple(tgt_tokens[-(no_repeat_ngram_size-1):])
+                
+                # Kiểm tra xem n-gram này đã xuất hiện trước đó chưa
+                for i in range(len(tgt_tokens) - no_repeat_ngram_size + 1):
+                    ngram = tuple(tgt_tokens[i:i+no_repeat_ngram_size-1])
+                    if ngram == last_ngram:
+                        # Nếu n-gram đã xuất hiện, giảm logit của token tiếp theo
+                        if i + no_repeat_ngram_size - 1 < len(tgt_tokens):
+                            repeated_token = tgt_tokens[i + no_repeat_ngram_size - 1]
+                            if repeated_token < len(next_token_logits):
+                                next_token_logits[repeated_token] = next_token_logits[repeated_token] / (repetition_penalty ** 2)
+            
+            # Lấy token có xác suất cao nhất
+            next_token = next_token_logits.argmax().item()
+            
+            # Thêm vào sequence
+            tgt_tokens.append(next_token)
+            
+            # Dừng nếu gặp <eos>
+            if next_token == tgt_vocab.EOS_IDX:
+                break
+        
+        # Decode thành sentence
+        decoded_sentence = tgt_vocab.decode(tgt_tokens)
+    
+    return tgt_tokens, decoded_sentence
+
+# ============================================================================
+# 2. BEAM SEARCH DECODING
+# ============================================================================
+
+class BeamSearchNode:
+    """
+    Node trong Beam Search
+    """
+    def __init__(self, tokens, log_prob, length):
+        self.tokens = tokens
+        self.log_prob = log_prob
+        self.length = length
+        
+    def eval(self, alpha=0.6):
+        """
+        Tính score với length normalization
+        
+        Score = log_prob / (length^alpha)
+        
+        Args:
+            alpha: Length penalty factor
+        """
+        return self.log_prob / (self.length ** alpha)
+
+def beam_search_decode(model, src, src_vocab, tgt_vocab, device, beam_size=5, max_len=100, alpha=0.6, repetition_penalty=1.5, no_repeat_ngram_size=3):
+    """
+    Beam Search Decoding - Giữ top-k candidates tốt nhất với Repetition Penalty
+    
+    Args:
+        model: Transformer model
+        src: Source sequence
+        src_vocab: Source vocabulary
+        tgt_vocab: Target vocabulary
+        device: Device
+        beam_size: Beam size (số lượng candidates)
+        max_len: Maximum length
+        alpha: Length penalty factor
+        repetition_penalty: Penalty factor cho tokens lặp lại (>1.0 để giảm repetition)
+        no_repeat_ngram_size: Kích thước n-gram để tránh lặp lại (0 = tắt)
+        
+    Returns:
+        best_tokens: List of best token indices
+        best_sentence: Best decoded sentence
+    """
+    model.eval()
+    
+    # Nếu src là list, convert sang tensor
+    if isinstance(src, list):
+        src = torch.LongTensor([src]).to(device)
+    elif src.dim() == 1:
+        src = src.unsqueeze(0)
+    
+    with torch.no_grad():
+        # Encode source
+        encoder_output, src_mask = model.encode(src)
+        
+        # Khởi tạo beam với <sos> token
+        beams = [BeamSearchNode(
+            tokens=[tgt_vocab.SOS_IDX],
+            log_prob=0.0,
+            length=1
+        )]
+        
+        completed_beams = []
+        
+        for step in range(max_len):
+            candidates = []
+            
+            for beam in beams:
+                # Nếu đã kết thúc, thêm vào completed
+                if beam.tokens[-1] == tgt_vocab.EOS_IDX:
+                    completed_beams.append(beam)
+                    continue
+                
+                # Tạo target tensor
+                tgt = torch.LongTensor([beam.tokens]).to(device)
+                
+                # Decode
+                output = model.decode(tgt, encoder_output, src_mask)
+                
+                # Lấy log probabilities cho token cuối
+                next_token_logits = output[0, -1, :]
+                
+                # Áp dụng repetition penalty
+                # Đếm số lần xuất hiện của mỗi token trong sequence hiện tại
+                token_counts = {}
+                for token_id in beam.tokens:
+                    if token_id not in [tgt_vocab.SOS_IDX, tgt_vocab.EOS_IDX, tgt_vocab.PAD_IDX]:
+                        token_counts[token_id] = token_counts.get(token_id, 0) + 1
+                
+                # Áp dụng penalty: giảm logit của các tokens đã xuất hiện (mạnh hơn)
+                for token_id, count in token_counts.items():
+                    if count > 0 and token_id < len(next_token_logits):
+                        # Penalty tăng mạnh theo số lần lặp lại
+                        # count=1: penalty nhẹ, count=2+: penalty rất mạnh
+                        penalty = repetition_penalty ** (count * 1.5)  # Tăng penalty mạnh hơn
+                        next_token_logits[token_id] = next_token_logits[token_id] / penalty
+                
+                # N-gram repetition penalty: tránh lặp cụm từ
+                if no_repeat_ngram_size > 0 and len(beam.tokens) >= no_repeat_ngram_size:
+                    # Lấy n-gram cuối cùng
+                    last_ngram = tuple(beam.tokens[-(no_repeat_ngram_size-1):])
+                    
+                    # Kiểm tra xem n-gram này đã xuất hiện trước đó chưa
+                    for i in range(len(beam.tokens) - no_repeat_ngram_size + 1):
+                        ngram = tuple(beam.tokens[i:i+no_repeat_ngram_size-1])
+                        if ngram == last_ngram:
+                            # Nếu n-gram đã xuất hiện, giảm logit của token tiếp theo rất mạnh
+                            if i + no_repeat_ngram_size - 1 < len(beam.tokens):
+                                repeated_token = beam.tokens[i + no_repeat_ngram_size - 1]
+                                if repeated_token < len(next_token_logits):
+                                    # Penalty rất mạnh cho n-gram repetition
+                                    next_token_logits[repeated_token] = next_token_logits[repeated_token] / (repetition_penalty ** 3)
+                
+                log_probs = F.log_softmax(next_token_logits, dim=-1)
+                
+                # Lấy top-k tokens
+                top_log_probs, top_tokens = torch.topk(log_probs, beam_size)
+                
+                # Tạo candidates mới
+                for log_prob, token in zip(top_log_probs, top_tokens):
+                    new_beam = BeamSearchNode(
+                        tokens=beam.tokens + [token.item()],
+                        log_prob=beam.log_prob + log_prob.item(),
+                        length=beam.length + 1
+                    )
+                    candidates.append(new_beam)
+            
+            # Nếu không còn candidates, dừng
+            if not candidates:
+                break
+            
+            # Chọn top beam_size candidates tốt nhất
+            beams = sorted(candidates, key=lambda x: x.eval(alpha), reverse=True)[:beam_size]
+            
+            # Nếu tất cả beams đã complete, dừng
+            if len(completed_beams) >= beam_size:
+                break
+        
+        # Thêm các beams chưa complete vào completed
+        completed_beams.extend(beams)
+        
+        # Chọn beam tốt nhất
+        best_beam = max(completed_beams, key=lambda x: x.eval(alpha))
+        
+        # Decode thành sentence
+        best_sentence = tgt_vocab.decode(best_beam.tokens)
+    
+    return best_beam.tokens, best_sentence
+
+# ============================================================================
+# 3. BLEU SCORE CALCULATION
+# ============================================================================
+
+def calculate_ngrams(tokens, n):
+    """
+    Tính n-grams từ list of tokens
+    
+    Args:
+        tokens: List of tokens
+        n: n-gram size
+        
+    Returns:
+        ngrams: Counter of n-grams
+    """
+    ngrams = []
+    for i in range(len(tokens) - n + 1):
+        ngram = tuple(tokens[i:i+n])
+        ngrams.append(ngram)
+    return Counter(ngrams)
+
+def calculate_bleu_score(references, hypotheses, max_n=4, weights=None):
+    """
+    Tính BLEU score
+    
+    BLEU = BP * exp(sum(w_n * log(p_n)))
+    
+    Args:
+        references: List of reference sentences (list of token lists)
+        hypotheses: List of hypothesis sentences (list of token lists)
+        max_n: Maximum n-gram size (mặc định 4)
+        weights: Weights cho mỗi n-gram (mặc định uniform)
+        
+    Returns:
+        bleu_score: BLEU score (0-100)
+    """
+    if weights is None:
+        weights = [1.0/max_n] * max_n
+    
+    # Tính precision cho mỗi n-gram
+    precisions = []
+    
+    for n in range(1, max_n + 1):
+        matched = 0
+        total = 0
+        
+        for ref, hyp in zip(references, hypotheses):
+            # Tính n-grams
+            ref_ngrams = calculate_ngrams(ref, n)
+            hyp_ngrams = calculate_ngrams(hyp, n)
+            
+            # Đếm matched n-grams
+            for ngram, count in hyp_ngrams.items():
+                matched += min(count, ref_ngrams.get(ngram, 0))
+            
+            total += max(len(hyp) - n + 1, 0)
+        
+        # Tính precision
+        if total > 0:
+            precision = matched / total
+        else:
+            precision = 0
+        
+        precisions.append(precision)
+    
+    # Tính Brevity Penalty (BP)
+    ref_length = sum(len(ref) for ref in references)
+    hyp_length = sum(len(hyp) for hyp in hypotheses)
+    
+    if hyp_length > ref_length:
+        bp = 1.0
+    elif hyp_length == 0:
+        bp = 0.0
+    else:
+        bp = math.exp(1 - ref_length / hyp_length)
+    
+    # Tính BLEU score
+    if min(precisions) > 0:
+        log_precisions = [w * math.log(p) for w, p in zip(weights, precisions)]
+        bleu = bp * math.exp(sum(log_precisions))
+    else:
+        bleu = 0.0
+    
+    return bleu * 100  # Convert sang 0-100 scale
+
+# ============================================================================
+# 4. EVALUATE ON TEST SET
+# ============================================================================
+
+def evaluate_model(model, test_loader, src_vocab, tgt_vocab, device, 
+                   use_beam_search=True, beam_size=5, max_len=100):
+    """
+    Đánh giá model trên test set
+    
+    Args:
+        model: Transformer model
+        test_loader: Test DataLoader
+        src_vocab: Source vocabulary
+        tgt_vocab: Target vocabulary
+        device: Device
+        use_beam_search: Sử dụng beam search hay greedy search
+        beam_size: Beam size (nếu dùng beam search)
+        max_len: Maximum decode length
+        
+    Returns:
+        bleu_score: BLEU score
+        translations: List of (source, reference, hypothesis) tuples
+    """
+    model.eval()
+    
+    references = []
+    hypotheses = []
+    translations = []
+    
+    print(f"\n{'='*70}")
+    print(f"ĐÁNH GIÁ TRÊN TEST SET")
+    print(f"{'='*70}")
+    print(f"Decoding method: {'Beam Search' if use_beam_search else 'Greedy Search'}")
+    if use_beam_search:
+        print(f"Beam size: {beam_size}")
+    print(f"{'='*70}\n")
+    
+    with torch.no_grad():
+        for src, tgt, _, _ in tqdm(test_loader, desc='Evaluating'):
+            src = src.to(device)
+            
+            batch_size = src.size(0)
+            
+            for i in range(batch_size):
+                src_seq = src[i]
+                tgt_seq = tgt[i]
+                
+                # Decode
+                if use_beam_search:
+                    _, hypothesis = beam_search_decode(
+                        model, src_seq, src_vocab, tgt_vocab, 
+                        device, beam_size, max_len, alpha=0.6, repetition_penalty=1.5, no_repeat_ngram_size=3
+                    )
+                else:
+                    _, hypothesis = greedy_decode(
+                        model, src_seq, src_vocab, tgt_vocab, 
+                        device, max_len, repetition_penalty=1.5, no_repeat_ngram_size=3
+                    )
+                
+                # Reference
+                reference = tgt_vocab.decode(tgt_seq.tolist())
+                
+                # Source
+                source = src_vocab.decode(src_seq.tolist())
+                
+                # Tokenize để tính BLEU
+                ref_tokens = reference.split()
+                hyp_tokens = hypothesis.split()
+                
+                references.append(ref_tokens)
+                hypotheses.append(hyp_tokens)
+                
+                translations.append((source, reference, hypothesis))
+    
+    # Tính BLEU score
+    bleu = calculate_bleu_score(references, hypotheses)
+    
+    print(f"\n{'='*70}")
+    print(f"KẾT QUẢ ĐÁNH GIÁ")
+    print(f"{'='*70}")
+    print(f"BLEU Score: {bleu:.2f}")
+    print(f"{'='*70}\n")
+    
+    return bleu, translations
+
+# ============================================================================
+# 5. PRINT SAMPLE TRANSLATIONS
+# ============================================================================
+
+def print_sample_translations(translations, num_samples=10):
+    """
+    In một số ví dụ dịch
+    
+    Args:
+        translations: List of (source, reference, hypothesis) tuples
+        num_samples: Số lượng samples để in
+    """
+    print(f"\n{'='*70}")
+    print(f"MỘT SỐ VÍ DỤ DỊCH")
+    print(f"{'='*70}\n")
+    
+    for i, (src, ref, hyp) in enumerate(translations[:num_samples], 1):
+        print(f"Ví dụ {i}:")
+        print(f"  Source:     {src}")
+        print(f"  Reference:  {ref}")
+        print(f"  Hypothesis: {hyp}")
+        print()
+
+# ============================================================================
+# 6. TRANSLATE SINGLE SENTENCE
+# ============================================================================
+
+def translate_sentence(model, sentence, src_vocab, tgt_vocab, device,
+                       use_beam_search=True, beam_size=5, max_len=100, src_lang='vi',
+                       repetition_penalty=1.5, no_repeat_ngram_size=3):
+    """
+    Dịch một câu đơn
+    
+    Args:
+        model: Transformer model
+        sentence: Source sentence (string)
+        src_vocab: Source vocabulary
+        tgt_vocab: Target vocabulary
+        device: Device
+        use_beam_search: Sử dụng beam search hay greedy
+        beam_size: Beam size
+        max_len: Maximum length
+        src_lang: Ngôn ngữ source ('vi' hoặc 'en')
+        
+    Returns:
+        translation: Translated sentence
+    """
+    model.eval()
+    
+    # Tiền xử lý câu
+    from data_preprocessing import clean_text
+    sentence = clean_text(sentence, src_lang)
+    
+    # Encode
+    tokens = src_vocab.encode(sentence)
+    src = torch.LongTensor([tokens]).to(device)
+    
+    # Decode
+    if use_beam_search:
+        _, translation = beam_search_decode(
+            model, src, src_vocab, tgt_vocab, 
+            device, beam_size, max_len, alpha=0.6, 
+            repetition_penalty=repetition_penalty, 
+            no_repeat_ngram_size=no_repeat_ngram_size
+        )
+    else:
+        _, translation = greedy_decode(
+            model, src, src_vocab, tgt_vocab, 
+            device, max_len, 
+            repetition_penalty=repetition_penalty, 
+            no_repeat_ngram_size=no_repeat_ngram_size
+        )
+    
+    return translation
+
+# ============================================================================
+# 7. INTERACTIVE TRANSLATION
+# ============================================================================
+
+def interactive_translation(model, src_vocab, tgt_vocab, device,
+                           use_beam_search=True, beam_size=5, direction='vi2en'):
+    """
+    Chế độ dịch tương tác
+    
+    Args:
+        model: Transformer model
+        src_vocab: Source vocabulary
+        tgt_vocab: Target vocabulary
+        device: Device
+        use_beam_search: Sử dụng beam search
+        beam_size: Beam size
+        direction: 'vi2en' hoặc 'en2vi'
+    """
+    print("\n" + "="*70)
+    print("CHẾ ĐỘ DỊCH TƯƠNG TÁC")
+    print("="*70)
+    if direction == 'vi2en':
+        print("Nhập câu tiếng Việt để dịch sang tiếng Anh")
+    else:
+        print("Nhập câu tiếng Anh để dịch sang tiếng Việt")
+    print("Gõ 'quit' hoặc 'exit' để thoát")
+    print("="*70 + "\n")
+    
+    src_lang = 'en' if direction == 'en2vi' else 'vi'
+    src_label = "Tiếng Anh" if direction == 'en2vi' else "Tiếng Việt"
+    tgt_label = "Tiếng Việt" if direction == 'en2vi' else "Tiếng Anh"
+    
+    while True:
+        sentence = input(f"{src_label}: ").strip()
+        
+        if sentence.lower() in ['quit', 'exit', '']:
+            print("Tạm biệt!")
+            break
+        
+        translation = translate_sentence(
+            model, sentence, src_vocab, tgt_vocab, device,
+            use_beam_search, beam_size, src_lang=src_lang, max_len=100
+        )
+        
+        print(f"{tgt_label}:  {translation}\n")
+
+# ============================================================================
+# 8. SAVE TRANSLATIONS
+# ============================================================================
+
+def save_translations(translations, output_file='translations.txt'):
+    """
+    Lưu translations ra file
+    
+    Args:
+        translations: List of (source, reference, hypothesis) tuples
+        output_file: Output file path
+    """
+    with open(output_file, 'w', encoding='utf-8') as f:
+        for i, (src, ref, hyp) in enumerate(translations, 1):
+            f.write(f"Example {i}:\n")
+            f.write(f"Source:     {src}\n")
+            f.write(f"Reference:  {ref}\n")
+            f.write(f"Hypothesis: {hyp}\n")
+            f.write("\n")
+    
+    print(f"✓ Saved translations to {output_file}")

src/inference_vlsp.py

@@ -0,0 +1,278 @@
+"""
+INFERENCE SCRIPT CHO VLSP
+Dịch test set và lưu kết quả vào file submission
+"""
+
+import torch
+import argparse
+from pathlib import Path
+from tqdm import tqdm
+import sys
+
+# Add parent directory to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from src.complete_transformer import TransformerShared
+from src.shared_vocab_utils import load_shared_tokenizer, load_shared_vocab_info
+from src.inference_evaluation import beam_search_decode, greedy_decode
+
+# ============================================================================
+# CONFIGURATION
+# ============================================================================
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DATA_DIR = PROJECT_ROOT / 'data'
+RESULTS_DIR = PROJECT_ROOT / 'results'
+CHECKPOINT_DIR = PROJECT_ROOT / 'checkpoints'
+
+# ============================================================================
+# LOAD MODEL
+# ============================================================================
+
+def load_model(checkpoint_path, model_size='custom_25m', device='cuda'):
+    """
+    Load model từ checkpoint
+    
+    Args:
+        checkpoint_path: Đường dẫn đến checkpoint file
+        model_size: Kích thước model ('custom_25m', 'base', etc.)
+        device: Device để load model
+        
+    Returns:
+        model: Loaded model
+    """
+    print(f"\n{'='*70}")
+    print(f"LOADING MODEL")
+    print(f"{'='*70}")
+    
+    # Load vocab info
+    vocab_info = load_shared_vocab_info()
+    vocab_size = vocab_info['vocab_size']
+    
+    print(f"Vocab size: {vocab_size}")
+    
+    # Model configs theo model_size
+    model_configs = {
+        'custom_25m': {
+            'd_model': 384,
+            'n_layers': 6,
+            'n_heads': 8,
+            'd_ff': 1536,
+            'dropout': 0.1
+        },
+        'base': {
+            'd_model': 512,
+            'n_layers': 6,
+            'n_heads': 8,
+            'd_ff': 2048,
+            'dropout': 0.1
+        }
+    }
+    
+    config = model_configs.get(model_size, model_configs['custom_25m'])
+    
+    # Tạo model
+    model = TransformerShared(
+        vocab_size=vocab_size,
+        d_model=config['d_model'],
+        n_layers=config['n_layers'],
+        n_heads=config['n_heads'],
+        d_ff=config['d_ff'],
+        dropout=config['dropout'],
+        pad_idx=vocab_info['pad_id']
+    )
+    
+    # Load checkpoint
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model = model.to(device)
+    model.eval()
+    
+    print(f"✓ Model loaded from {checkpoint_path}")
+    print(f"✓ Model size: {model_size}")
+    print(f"✓ Device: {device}")
+    print(f"{'='*70}\n")
+    
+    return model, vocab_info
+
+# ============================================================================
+# VOCABULARY WRAPPER
+# ============================================================================
+
+class VocabWrapper:
+    """Wrapper cho tokenizer để tương thích với inference functions"""
+    def __init__(self, tokenizer, vocab_info):
+        self.tokenizer = tokenizer
+        self.SOS_IDX = vocab_info['sos_id']
+        self.EOS_IDX = vocab_info['eos_id']
+        self.PAD_IDX = vocab_info['pad_id']
+        self.UNK_IDX = vocab_info['unk_id']
+    
+    def encode(self, text):
+        """Encode text thành token IDs"""
+        encoding = self.tokenizer.encode(text)
+        return encoding.ids
+    
+    def decode(self, token_ids):
+        """Decode token IDs thành text (nhận list of token IDs)"""
+        # Filter out special tokens
+        filtered_ids = [tid for tid in token_ids 
+                       if tid not in [self.SOS_IDX, self.EOS_IDX, self.PAD_IDX]]
+        if not filtered_ids:
+            return ""
+        # Decode using tokenizer
+        decoded = self.tokenizer.decode(filtered_ids)
+        return decoded
+
+# ============================================================================
+# INFERENCE
+# ============================================================================
+
+def translate_file(
+    model, 
+    tokenizer, 
+    vocab_info,
+    source_file,
+    output_file,
+    device='cuda',
+    use_beam_search=True,
+    beam_size=5,
+    max_len=100
+):
+    """
+    Dịch file source và lưu kết quả
+    
+    Args:
+        model: Transformer model
+        tokenizer: Tokenizer
+        vocab_info: Vocabulary info dict
+        source_file: File chứa source sentences (mỗi dòng 1 câu)
+        output_file: File output để lưu translations
+        device: Device
+        use_beam_search: Dùng beam search hay greedy
+        beam_size: Beam size
+        max_len: Max decode length
+    """
+    print(f"\n{'='*70}")
+    print(f"INFERENCE")
+    print(f"{'='*70}")
+    print(f"Source file: {source_file}")
+    print(f"Output file: {output_file}")
+    print(f"Method: {'Beam Search' if use_beam_search else 'Greedy Search'}")
+    if use_beam_search:
+        print(f"Beam size: {beam_size}")
+    print(f"{'='*70}\n")
+    
+    # Load source sentences
+    with open(source_file, 'r', encoding='utf-8') as f:
+        source_sentences = [line.strip() for line in f.readlines()]
+    
+    print(f"Loaded {len(source_sentences)} source sentences")
+    
+    # Create vocab wrapper
+    vocab = VocabWrapper(tokenizer, vocab_info)
+    
+    # Translate
+    translations = []
+    model.eval()
+    
+    with torch.no_grad():
+        for src_text in tqdm(source_sentences, desc='Translating'):
+            # Encode source
+            src_ids = vocab.encode(src_text)
+            src_tensor = torch.LongTensor([src_ids]).to(device)
+            
+            # Decode
+            if use_beam_search:
+                _, translation = beam_search_decode(
+                    model, src_tensor, vocab, vocab, device, beam_size, max_len, 
+                    alpha=0.6, repetition_penalty=1.5, no_repeat_ngram_size=3
+                )
+            else:
+                _, translation = greedy_decode(
+                    model, src_tensor, vocab, vocab, device, max_len,
+                    repetition_penalty=1.5, no_repeat_ngram_size=3
+                )
+            
+            translations.append(translation)
+    
+    # Save translations
+    RESULTS_DIR.mkdir(exist_ok=True, parents=True)
+    with open(output_file, 'w', encoding='utf-8') as f:
+        for trans in translations:
+            f.write(trans + '\n')
+    
+    print(f"\n✓ Saved {len(translations)} translations to {output_file}")
+    print(f"{'='*70}\n")
+
+# ============================================================================
+# MAIN
+# ============================================================================
+
+def main():
+    parser = argparse.ArgumentParser(description='Inference cho VLSP')
+    parser.add_argument('--checkpoint', type=str, required=True,
+                       help='Đường dẫn đến checkpoint file')
+    parser.add_argument('--source', type=str, 
+                       default='data/raw/Hospital/public_test.en.txt',
+                       help='File source (tiếng Anh)')
+    parser.add_argument('--output', type=str,
+                       default='results/submission_vlsp.txt',
+                       help='File output (tiếng Việt)')
+    parser.add_argument('--model_size', type=str, default='custom_25m',
+                       choices=['custom_25m', 'base'],
+                       help='Kích thước model')
+    parser.add_argument('--beam_size', type=int, default=5,
+                       help='Beam size cho beam search')
+    parser.add_argument('--greedy', action='store_true',
+                       help='Dùng greedy search thay vì beam search')
+    parser.add_argument('--max_len', type=int, default=100,
+                       help='Maximum decode length')
+    parser.add_argument('--device', type=str, default='cuda',
+                       help='Device (cuda/cpu)')
+    
+    args = parser.parse_args()
+    
+    # Check files
+    checkpoint_path = Path(args.checkpoint)
+    if not checkpoint_path.exists():
+        print(f"❌ Không tìm thấy checkpoint: {checkpoint_path}")
+        return
+    
+    source_path = Path(args.source)
+    if not source_path.exists():
+        print(f"❌ Không tìm thấy source file: {source_path}")
+        return
+    
+    # Device
+    device = args.device
+    if device == 'cuda' and not torch.cuda.is_available():
+        print("⚠️  CUDA không khả dụng, dùng CPU")
+        device = 'cpu'
+    
+    # Load model
+    model, vocab_info = load_model(
+        checkpoint_path, 
+        args.model_size, 
+        device
+    )
+    
+    # Load tokenizer
+    tokenizer = load_shared_tokenizer()
+    
+    # Translate
+    translate_file(
+        model=model,
+        tokenizer=tokenizer,
+        vocab_info=vocab_info,
+        source_file=source_path,
+        output_file=Path(args.output),
+        device=device,
+        use_beam_search=not args.greedy,
+        beam_size=args.beam_size,
+        max_len=args.max_len
+    )
+
+if __name__ == '__main__':
+    main()
+

src/main_pipeline.py

@@ -0,0 +1,550 @@
+"""
+MAIN PIPELINE - HOÀN CHỈNH (FIXED)
+Script chính để chạy toàn bộ pipeline: Data → Train → Evaluate
+"""
+
+import torch
+import argparse
+import os
+import pickle
+import json
+from pathlib import Path
+
+# Import modules - IMPORTANT: Import Vocabulary class
+from data_preprocessing import main as prepare_data, Vocabulary
+from dataloader_module import create_dataloaders_with_bucketing, load_data_and_vocab
+from complete_transformer import create_model, print_model_info
+from training_module import train_model, plot_training_history, load_checkpoint
+from inference_evaluation import (
+    evaluate_model, print_sample_translations, 
+    save_translations, interactive_translation
+)
+
+# ============================================================================
+# PATH SETUP & CONFIGURATION
+# ============================================================================
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DATA_DIR = PROJECT_ROOT / 'data'
+PROCESSED_DATA_DIR = DATA_DIR / 'processed'
+CHECKPOINT_DIR = PROJECT_ROOT / 'checkpoints'
+RESULTS_DIR = PROJECT_ROOT / 'results'
+LOGS_DIR = PROJECT_ROOT / 'logs'
+
+
+class Config:
+    """
+    Configuration cho training
+    """
+    # Data
+    data_dir = DATA_DIR
+    checkpoint_dir = CHECKPOINT_DIR
+    results_dir = RESULTS_DIR
+    logs_dir = LOGS_DIR
+    
+    # Model
+    # Cấu hình cho 1M dataset: ~25M parameters, 32k shared vocab
+    model_size = 'medium'  # 'tiny', 'small', 'medium' (~25M), 'base', 'large'
+    use_weight_tying = True   # Weight tying enabled (tiết kiệm ~50% vocab params)
+    
+    # Training - Tối ưu cho RTX 5060 Ti 16GB
+    num_epochs = 20
+    batch_size = 128  # Batch size lớn để tận dụng GPU
+    grad_accum_steps = 2  # Effective batch = 128 * 2 = 256
+    warmup_steps = 4000
+    label_smoothing = 0.1
+    gradient_clip = 1.0
+    precision = 'bf16'  # 'bf16' (recommended), 'amp' (float16), hoặc 'fp32'
+    
+    # Decoding
+    beam_size = 5
+    max_decode_len = 100
+    
+    # Device
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
+    # Num workers - set to 0 on Windows to avoid multiprocessing issues
+    import platform
+    num_workers = 0 if platform.system() == 'Windows' else 2
+    
+    # Checkpoint
+    save_every = 1
+    
+    def __str__(self):
+        config_str = "\n" + "="*70 + "\n"
+        config_str += "CONFIGURATION\n"
+        config_str += "="*70 + "\n"
+        for key, value in self.__dict__.items():
+            if not key.startswith('_'):
+                config_str += f"{key}: {value}\n"
+        config_str += "="*70 + "\n"
+        return config_str
+
+# ============================================================================
+# MAIN FUNCTIONS
+# ============================================================================
+
+def setup_directories():
+    """
+    Tạo các thư mục cần thiết
+    """
+    for directory in [CHECKPOINT_DIR, RESULTS_DIR, LOGS_DIR]:
+        os.makedirs(directory, exist_ok=True)
+    print("✓ Created directories")
+
+def stage_1_prepare_data():
+    """
+    STAGE 1: Chuẩn bị dữ liệu (Shared Vocabulary) - 1M dataset từ MTet
+    """
+    print("\n" + "="*70)
+    print("STAGE 1: CHUẨN BỊ DỮ LIỆU (SHARED VOCABULARY - 1M DATASET)")
+    print("="*70 + "\n")
+    
+    # Kiểm tra xem đã có shared vocab data chưa
+    shared_data_files = [
+        PROCESSED_DATA_DIR / 'tokenizer_shared.json',
+        PROCESSED_DATA_DIR / 'processed_data_shared.pkl',
+        PROCESSED_DATA_DIR / 'shared_vocab_info.json'
+    ]
+    
+    if all(os.path.exists(f) for f in shared_data_files):
+        print("✓ Shared vocabulary data đã được chuẩn bị trước đó")
+        print(f"  - Tokenizer: {shared_data_files[0]}")
+        print(f"  - Processed data: {shared_data_files[1]}")
+        print(f"  - Vocab info: {shared_data_files[2]}")
+        
+        # Kiểm tra số lượng data
+        import pickle
+        with open(shared_data_files[1], 'rb') as f:
+            data = pickle.load(f)
+        train_size = len(data.get('train', []))
+        print(f"  - Train samples: {train_size:,}")
+        
+        if train_size < 500000:
+            print(f"\n⚠️  Dataset hiện tại chỉ có {train_size:,} câu")
+            print("   Để train 1M dataset, cần tạo lại từ MTet CSV")
+            response = input("Bạn có muốn tạo lại với 1M câu từ MTet không? (y/n): ")
+            if response.lower() == 'y':
+                print("\n→ Đang tạo lại với 1M câu từ MTet...")
+                import os
+                os.environ['MTET_MAX_ROWS'] = '1000000'
+                # Chạy lại build vocab và encode
+                from subprocess import run
+                run(['python3', 'src/1_build_shared_vocab.py'], check=True)
+                run(['python3', 'src/2_encode_data.py'], check=True)
+                print("✓ Đã tạo lại với 1M câu!")
+        else:
+            response = input("Bạn có mu���n chuẩn bị lại không? (y/n): ")
+            if response.lower() != 'y':
+                print("→ Bỏ qua stage 1")
+                return
+    else:
+        # Tự động setup với 1M câu từ MTet
+        print("📝 Tự động setup shared vocabulary với 1M câu từ MTet...")
+        import os
+        os.environ['MTET_MAX_ROWS'] = '1000000'
+        
+        print("\n→ Chạy: python src/1_build_shared_vocab.py")
+        from subprocess import run
+        try:
+            run(['python3', 'src/1_build_shared_vocab.py'], check=True)
+            print("\n→ Chạy: python src/2_encode_data.py")
+            run(['python3', 'src/2_encode_data.py'], check=True)
+            print("\n✓ Đã setup xong với 1M câu!")
+        except Exception as e:
+            print(f"\n⚠️  Lỗi khi setup: {e}")
+            print("Vui lòng chạy thủ công:")
+            print("  1. export MTET_MAX_ROWS=1000000")
+            print("  2. python src/1_build_shared_vocab.py")
+            print("  3. python src/2_encode_data.py")
+            return
+    
+    print("\n✓ HOÀN TẤT STAGE 1")
+
+def stage_2_create_model(config):
+    """
+    STAGE 2: Tạo model với Shared Vocabulary
+    """
+    print("\n" + "="*70)
+    print("STAGE 2: TẠO MODEL (SHARED VOCABULARY)")
+    print("="*70 + "\n")
+    
+    # Load shared vocabulary
+    from shared_vocab_utils import load_shared_vocab_info, create_shared_vocab_wrapper
+    
+    vocab_info = load_shared_vocab_info()
+    vocab_size = vocab_info['vocab_size']
+    pad_idx = vocab_info['pad_id']
+    
+    print(f"Loading shared vocabulary:")
+    print(f"  - Vocab size: {vocab_size}")
+    print(f"  - PAD ID: {pad_idx}")
+    print(f"  - SOS ID: {vocab_info['sos_id']}")
+    print(f"  - EOS ID: {vocab_info['eos_id']}")
+    
+    # Tạo shared vocab wrappers
+    vi_vocab, en_vocab = create_shared_vocab_wrapper()
+    
+    print(f"\n✓ Shared vocabulary loaded (cả 2 ngôn ngữ dùng chung)")
+    
+    # Tạo model với shared vocab
+    model, model_config = create_model(
+        src_vocab_size=vocab_size,
+        tgt_vocab_size=vocab_size,
+        model_size=config.model_size,
+        pad_idx=pad_idx,
+        use_shared_vocab=True,
+        use_weight_tying=config.use_weight_tying
+    )
+    
+    # Print model info
+    print_model_info(model, config.model_size, use_shared_vocab=True)
+    
+    # Move to device
+    device = torch.device(config.device)
+    model = model.to(device)
+    
+    # Kiểm tra và load checkpoint cũ nếu có (chỉ khi KHÔNG train bidirectional)
+    # Nếu train bidirectional, không load checkpoint cũ để train từ đầu
+    use_bidirectional = getattr(config, 'use_bidirectional', False)
+    
+    if not use_bidirectional:
+        # Chỉ load checkpoint khi train vi-en (không bidirectional)
+        checkpoint_dir = Path(config.checkpoint_dir)
+        best_model_path = checkpoint_dir / 'best_model.pt'
+        
+        # Tìm checkpoint epoch cao nhất
+        max_epoch = 0
+        latest_checkpoint = None
+        for checkpoint_file in checkpoint_dir.glob('checkpoint_epoch_*.pt'):
+            try:
+                epoch_num = int(checkpoint_file.stem.split('_')[-1])
+                if epoch_num > max_epoch:
+                    max_epoch = epoch_num
+                    latest_checkpoint = checkpoint_file
+            except:
+                pass
+        
+        # Load checkpoint nếu có
+        if latest_checkpoint and latest_checkpoint.exists():
+            print(f"\n→ Tìm thấy checkpoint cũ: {latest_checkpoint.name} (epoch {max_epoch})")
+            print(f"→ Model sẽ được load từ checkpoint này để tiếp tục training")
+            try:
+                checkpoint = torch.load(latest_checkpoint, map_location=device)
+                model.load_state_dict(checkpoint['model_state_dict'])
+                print(f"✓ Đã load weights từ checkpoint epoch {max_epoch}")
+            except Exception as e:
+                print(f"⚠️  Không thể load checkpoint: {e}")
+                print(f"   Sẽ train từ đầu với weights mới")
+        elif best_model_path.exists():
+            print(f"\n→ Tìm thấy best_model.pt")
+            print(f"→ Model sẽ được load từ best_model để tiếp tục training")
+            try:
+                checkpoint = torch.load(best_model_path, map_location=device)
+                model.load_state_dict(checkpoint['model_state_dict'])
+                print(f"✓ Đã load weights từ best_model")
+            except Exception as e:
+                print(f"⚠️  Không thể load checkpoint: {e}")
+                print(f"   Sẽ train từ đầu với weights mới")
+        else:
+            print(f"\n→ Không tìm thấy checkpoint cũ")
+            print(f"→ Sẽ train từ đầu với weights mới")
+    else:
+        # Train bidirectional: KHÔNG load checkpoint cũ, train từ đầu
+        print(f"\n→ Training với bidirectional dataset")
+        print(f"→ KHÔNG load checkpoint cũ, sẽ train từ đầu với weights mới")
+        print(f"→ Dropout: {model_config.get('dropout', 0.1)} (đã có trong model)")
+    
+    print(f"\n✓ Model moved to {device}")
+    print("✓ HOÀN TẤT STAGE 2")
+    
+    return model, vi_vocab, en_vocab, model_config
+
+def stage_3_prepare_dataloaders(config):
+    """
+    STAGE 3: Chuẩn bị DataLoaders với Shared Vocabulary
+    """
+    print("\n" + "="*70)
+    print("STAGE 3: CHUẨN BỊ DATALOADERS (SHARED VOCABULARY)")
+    print("="*70 + "\n")
+    
+    # Load shared vocab data
+    from dataloader_module import load_data_and_vocab
+    
+    # Check if bidirectional flag is set
+    use_bidirectional = getattr(config, 'use_bidirectional', False)
+    
+    processed_data, vi_vocab, en_vocab = load_data_and_vocab(
+        use_shared_vocab=True,
+        use_bidirectional=use_bidirectional
+    )
+    
+    print(f"✓ Loaded shared vocab data:")
+    print(f"  - Train: {len(processed_data['train']['src'])} cặp câu")
+    print(f"  - Validation: {len(processed_data['validation']['src'])} cặp câu")
+    print(f"  - Test: {len(processed_data['test']['src'])} cặp câu")
+    
+    # Tạo dataloaders
+    train_loader, val_loader, test_loader = create_dataloaders_with_bucketing(
+        processed_data,
+        batch_size=config.batch_size,
+        num_workers=config.num_workers
+    )
+    
+    print("\n✓ HOÀN TẤT STAGE 3")
+    
+    return train_loader, val_loader, test_loader
+
+def stage_4_train_model(model, train_loader, val_loader, config, model_config):
+    """
+    STAGE 4: Huấn luyện model
+    """
+    print("\n" + "="*70)
+    print("STAGE 4: HUẤN LUYỆN MODEL")
+    print("="*70 + "\n")
+    
+    # Kiểm tra có checkpoint không và tự động tiếp tục
+    # CHỈ resume khi KHÔNG train bidirectional
+    use_bidirectional = getattr(config, 'use_bidirectional', False)
+    resume_from = None
+    
+    if not use_bidirectional:
+        # Chỉ resume khi train vi-en (không bidirectional)
+        checkpoint_dir = Path(config.checkpoint_dir)
+        
+        # Tìm checkpoint epoch cao nhất
+        max_epoch = 0
+        for checkpoint_file in checkpoint_dir.glob('checkpoint_epoch_*.pt'):
+            try:
+                epoch_num = int(checkpoint_file.stem.split('_')[-1])
+                if epoch_num > max_epoch:
+                    max_epoch = epoch_num
+                    resume_from = checkpoint_file
+            except:
+                pass
+        
+        # Nếu không tìm thấy checkpoint epoch, thử dùng best_model
+        if resume_from is None or not resume_from.exists():
+            best_model_path = checkpoint_dir / 'best_model.pt'
+            if best_model_path.exists():
+                # Kiểm tra epoch trong best_model
+                try:
+                    checkpoint = torch.load(best_model_path, map_location='cpu')
+                    if 'epoch' in checkpoint:
+                        max_epoch = checkpoint['epoch']
+                        resume_from = best_model_path
+                        print(f"→ Tìm thấy best_model.pt (epoch {max_epoch})")
+                        print(f"→ Sẽ tiếp tục training từ epoch {max_epoch + 1}")
+                except:
+                    pass
+        
+        if resume_from and resume_from.exists():
+            print(f"→ Tìm thấy checkpoint: {resume_from.name}")
+            print(f"→ Tự động tiếp tục training từ epoch {max_epoch + 1}")
+        else:
+            print(f"→ Không tìm thấy checkpoint cũ, sẽ train từ đầu")
+    else:
+        # Train bidirectional: KHÔNG resume, train từ đầu
+        print(f"→ Training với bidirectional dataset")
+        print(f"→ KHÔNG resume từ checkpoint cũ, sẽ train từ đầu (epoch 1)")
+        resume_from = None
+    
+    # Train
+    history = train_model(
+        model=model,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        num_epochs=config.num_epochs,
+        device=torch.device(config.device),
+        d_model=model_config['d_model'],
+        warmup_steps=config.warmup_steps,
+        label_smoothing=config.label_smoothing,
+        grad_accum_steps=config.grad_accum_steps,
+        precision=config.precision,
+        checkpoint_dir=config.checkpoint_dir,
+        save_every=config.save_every,
+        resume_from=str(resume_from) if resume_from else None
+    )
+    
+    # Plot history
+    training_plot_path = Path(config.results_dir) / 'training_history.png'
+    plot_training_history(history, save_path=training_plot_path)
+    
+    # Save history
+    training_history_path = Path(config.results_dir) / 'training_history.json'
+    with open(training_history_path, 'w') as f:
+        json.dump(history, f, indent=2)
+    
+    print("\n✓ HOÀN TẤT STAGE 4")
+    
+    return history
+
+def stage_5_evaluate_model(model, test_loader, vi_vocab, en_vocab, config):
+    """
+    STAGE 5: Đánh giá model
+    """
+    print("\n" + "="*70)
+    print("STAGE 5: ĐÁNH GIÁ MODEL")
+    print("="*70 + "\n")
+    
+    device = torch.device(config.device)
+    
+    # Load best model
+    best_model_path = Path(config.checkpoint_dir) / 'best_model.pt'
+    if os.path.exists(best_model_path):
+        print("Loading best model...")
+        model, _ = load_checkpoint(model, best_model_path, device)
+    
+    # Evaluate với Greedy Search
+    print("\n--- GREEDY SEARCH ---")
+    bleu_greedy, translations_greedy = evaluate_model(
+        model, test_loader, vi_vocab, en_vocab, device,
+        use_beam_search=False, max_len=config.max_decode_len
+    )
+    
+    # Evaluate với Beam Search
+    print("\n--- BEAM SEARCH ---")
+    bleu_beam, translations_beam = evaluate_model(
+        model, test_loader, vi_vocab, en_vocab, device,
+        use_beam_search=True, beam_size=config.beam_size, max_len=config.max_decode_len
+    )
+    
+    # In samples
+    print("\n" + "="*70)
+    print("GREEDY SEARCH SAMPLES")
+    print("="*70)
+    print_sample_translations(translations_greedy, num_samples=10)
+    
+    print("\n" + "="*70)
+    print("BEAM SEARCH SAMPLES")
+    print("="*70)
+    print_sample_translations(translations_beam, num_samples=10)
+    
+    # Save translations
+    save_translations(translations_greedy, Path(config.results_dir) / 'translations_greedy.txt')
+    save_translations(translations_beam, Path(config.results_dir) / 'translations_beam.txt')
+    
+    # Save scores
+    scores = {
+        'greedy_bleu': bleu_greedy,
+        'beam_bleu': bleu_beam,
+        'beam_size': config.beam_size
+    }
+    
+    with open(Path(config.results_dir) / 'bleu_scores.json', 'w') as f:
+        json.dump(scores, f, indent=2)
+    
+    print("\n" + "="*70)
+    print("KẾT QUẢ CUỐI CÙNG")
+    print("="*70)
+    print(f"Greedy Search BLEU: {bleu_greedy:.2f}")
+    print(f"Beam Search BLEU: {bleu_beam:.2f}")
+    print(f"Improvement: {bleu_beam - bleu_greedy:.2f}")
+    print("="*70)
+    
+    print("\n✓ HOÀN TẤT STAGE 5")
+    
+    return bleu_greedy, bleu_beam
+
+def stage_6_interactive_mode(model, vi_vocab, en_vocab, config):
+    """
+    STAGE 6: Chế độ dịch tương tác
+    """
+    print("\n" + "="*70)
+    print("STAGE 6: CHẾ ĐỘ DỊCH TƯƠNG TÁC")
+    print("="*70 + "\n")
+    
+    device = torch.device(config.device)
+    
+    # Load best model
+    best_model_path = os.path.join(config.checkpoint_dir, 'best_model.pt')
+    if os.path.exists(best_model_path):
+        model, _ = load_checkpoint(model, best_model_path, device)
+    
+    # Interactive translation
+    interactive_translation(
+        model, vi_vocab, en_vocab, device,
+        use_beam_search=True, beam_size=config.beam_size
+    )
+    
+    print("\n✓ HOÀN TẤT STAGE 6")
+
+# ============================================================================
+# MAIN PIPELINE
+# ============================================================================
+
+def main():
+    """
+    Main pipeline
+    """
+    # Parse arguments
+    parser = argparse.ArgumentParser(description='Transformer Machine Translation')
+    parser.add_argument('--stage', type=str, default='all', 
+                       help='Which stage to run: all, 1, 2, 3, 4, 5, 6')
+    parser.add_argument('--model_size', type=str, default='medium',
+                       help='Model size: tiny, small, medium (~25M), base, large')
+    parser.add_argument('--epochs', type=int, default=20,
+                       help='Number of epochs')
+    parser.add_argument('--batch_size', type=int, default=32,
+                       help='Batch size')
+    parser.add_argument('--beam_size', type=int, default=5,
+                       help='Beam size for beam search')
+    parser.add_argument('--grad_accum_steps', type=int, default=2,
+                       help='Gradient accumulation steps (effective batch = batch_size * grad_accum_steps)')
+    parser.add_argument('--precision', type=str, default='bf16', choices=['bf16', 'amp', 'fp32'],
+                       help='Training precision: bf16 (bfloat16, recommended), amp (float16), or fp32')
+    parser.add_argument('--no_weight_tying', action='store_true',
+                       help='Disable weight tying (default: weight tying enabled)')
+    parser.add_argument('--bidirectional', action='store_true',
+                       help='Use bidirectional dataset (vi↔en)')
+    args = parser.parse_args()
+    
+    # Setup config
+    config = Config()
+    config.model_size = args.model_size
+    config.num_epochs = args.epochs
+    config.use_bidirectional = args.bidirectional
+    config.batch_size = args.batch_size
+    config.beam_size = args.beam_size
+    config.grad_accum_steps = max(1, args.grad_accum_steps)
+    config.precision = args.precision
+    config.use_weight_tying = not args.no_weight_tying
+    
+    print(config)
+    
+    # Setup directories
+    setup_directories()
+    
+    # Run stages
+    stage = args.stage.lower()
+    
+    if stage in ['all', '1']:
+        stage_1_prepare_data()
+    
+    if stage in ['all', '2', '3', '4', '5', '6']:
+        model, vi_vocab, en_vocab, model_config = stage_2_create_model(config)
+    
+    if stage in ['all', '3', '4', '5']:
+        train_loader, val_loader, test_loader = stage_3_prepare_dataloaders(config)
+    
+    if stage in ['all', '4']:
+        history = stage_4_train_model(model, train_loader, val_loader, config, model_config)
+    
+    if stage in ['all', '5']:
+        bleu_greedy, bleu_beam = stage_5_evaluate_model(
+            model, test_loader, vi_vocab, en_vocab, config
+        )
+    
+    if stage in ['6', 'interactive']:
+        stage_6_interactive_mode(model, vi_vocab, en_vocab, config)
+    
+    print("\n" + "="*70)
+    print("🎉 HOÀN TẤT TOÀN BỘ PIPELINE!")
+    print("="*70)
+
+# ============================================================================
+# RUN
+# ============================================================================
+
+if __name__ == "__main__":
+    main()

src/prepare_hospital_data.py

@@ -0,0 +1,204 @@
+"""
+PREPARE HOSPITAL DATA FOR FINETUNING
+Xử lý dữ liệu Hospital domain để finetune model
+"""
+
+import pickle
+import json
+from pathlib import Path
+from tqdm import tqdm
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DATA_DIR = PROJECT_ROOT / 'data'
+RAW_DATA_DIR = DATA_DIR / 'raw'
+PROCESSED_DATA_DIR = DATA_DIR / 'processed'
+HOSPITAL_DIR = RAW_DATA_DIR / 'Hospital'
+
+def load_hospital_data():
+    """
+    Load dữ liệu Hospital từ các file .txt
+    """
+    print("="*70)
+    print("LOADING HOSPITAL DATA")
+    print("="*70)
+    
+    train_vi_path = HOSPITAL_DIR / 'train.vi.txt'
+    train_en_path = HOSPITAL_DIR / 'train.en.txt'
+    test_vi_path = HOSPITAL_DIR / 'public_test.vi.txt'
+    test_en_path = HOSPITAL_DIR / 'public_test.en.txt'
+    
+    # Load training data
+    print("\n→ Loading training data...")
+    with open(train_vi_path, 'r', encoding='utf-8') as f:
+        train_vi = [line.strip() for line in f if line.strip()]
+    with open(train_en_path, 'r', encoding='utf-8') as f:
+        train_en = [line.strip() for line in f if line.strip()]
+    
+    print(f"  ✓ Train VI: {len(train_vi)} câu")
+    print(f"  ✓ Train EN: {len(train_en)} câu")
+    
+    # Load test data
+    print("\n→ Loading test data...")
+    with open(test_vi_path, 'r', encoding='utf-8') as f:
+        test_vi = [line.strip() for line in f if line.strip()]
+    with open(test_en_path, 'r', encoding='utf-8') as f:
+        test_en = [line.strip() for line in f if line.strip()]
+    
+    print(f"  ✓ Test VI: {len(test_vi)} câu")
+    print(f"  ✓ Test EN: {len(test_en)} câu")
+    
+    # Kiểm tra số lượng khớp
+    if len(train_vi) != len(train_en):
+        print(f"⚠️  Warning: Train VI ({len(train_vi)}) != Train EN ({len(train_en)})")
+        min_len = min(len(train_vi), len(train_en))
+        train_vi = train_vi[:min_len]
+        train_en = train_en[:min_len]
+        print(f"  → Đã cắt về {min_len} cặp câu")
+    
+    if len(test_vi) != len(test_en):
+        print(f"⚠️  Warning: Test VI ({len(test_vi)}) != Test EN ({len(test_en)})")
+        min_len = min(len(test_vi), len(test_en))
+        test_vi = test_vi[:min_len]
+        test_en = test_en[:min_len]
+        print(f"  → Đã cắt về {min_len} cặp câu")
+    
+    # Split train thành train và validation (90-10)
+    print("\n→ Splitting train/validation (90-10)...")
+    split_idx = int(len(train_vi) * 0.9)
+    val_vi = train_vi[split_idx:]
+    val_en = train_en[split_idx:]
+    train_vi = train_vi[:split_idx]
+    train_en = train_en[:split_idx]
+    
+    print(f"  ✓ Train: {len(train_vi)} cặp câu")
+    print(f"  ✓ Validation: {len(val_vi)} cặp câu")
+    print(f"  ✓ Test: {len(test_vi)} cặp câu")
+    
+    return {
+        'train': {'vi': train_vi, 'en': train_en},
+        'validation': {'vi': val_vi, 'en': val_en},
+        'test': {'vi': test_vi, 'en': test_en}
+    }
+
+def encode_hospital_data(hospital_data):
+    """
+    Encode dữ liệu Hospital với shared vocabulary hiện có
+    """
+    print("\n" + "="*70)
+    print("ENCODING HOSPITAL DATA")
+    print("="*70)
+    
+    # Load shared vocabulary
+    from shared_vocab_utils import load_shared_vocab_info, create_shared_vocab_wrapper
+    
+    print("\n→ Loading shared vocabulary...")
+    vocab_info = load_shared_vocab_info()
+    vi_vocab, en_vocab = create_shared_vocab_wrapper()
+    
+    print(f"  ✓ Vocab size: {vocab_info['vocab_size']}")
+    print(f"  ✓ PAD ID: {vocab_info['pad_id']}")
+    print(f"  ✓ SOS ID: {vocab_info['sos_id']}")
+    print(f"  ✓ EOS ID: {vocab_info['eos_id']}")
+    
+    # Encode data
+    print("\n→ Encoding data...")
+    encoded_data = {}
+    
+    for split in ['train', 'validation', 'test']:
+        print(f"\n  Encoding {split}...")
+        vi_sentences = hospital_data[split]['vi']
+        en_sentences = hospital_data[split]['en']
+        
+        encoded_vi = []
+        encoded_en = []
+        
+        for vi_sent, en_sent in tqdm(zip(vi_sentences, en_sentences), 
+                                     total=len(vi_sentences), 
+                                     desc=f"  {split}"):
+            # Encode VI
+            vi_tokens = vi_vocab.encode(vi_sent)
+            encoded_vi.append(vi_tokens)
+            
+            # Encode EN
+            en_tokens = en_vocab.encode(en_sent)
+            encoded_en.append(en_tokens)
+        
+        encoded_data[split] = {
+            'src': encoded_vi,
+            'tgt': encoded_en
+        }
+        
+        print(f"    ✓ {split}: {len(encoded_vi)} cặp câu")
+    
+    return encoded_data, vi_vocab, en_vocab
+
+def save_hospital_data(encoded_data):
+    """
+    Lưu dữ liệu Hospital đã encode
+    """
+    print("\n" + "="*70)
+    print("SAVING HOSPITAL DATA")
+    print("="*70)
+    
+    hospital_data_path = PROCESSED_DATA_DIR / 'hospital_data_encoded.pkl'
+    
+    print(f"\n→ Saving to {hospital_data_path}...")
+    with open(hospital_data_path, 'wb') as f:
+        pickle.dump(encoded_data, f)
+    
+    print(f"  ✓ Saved: {hospital_data_path}")
+    
+    # Save info
+    info = {
+        'train_size': len(encoded_data['train']['src']),
+        'validation_size': len(encoded_data['validation']['src']),
+        'test_size': len(encoded_data['test']['src']),
+        'total_size': sum(len(encoded_data[split]['src']) for split in ['train', 'validation', 'test'])
+    }
+    
+    info_path = PROCESSED_DATA_DIR / 'hospital_data_info.json'
+    with open(info_path, 'w', encoding='utf-8') as f:
+        json.dump(info, f, indent=2, ensure_ascii=False)
+    
+    print(f"  ✓ Saved info: {info_path}")
+    print(f"\n  Summary:")
+    print(f"    - Train: {info['train_size']:,} cặp câu")
+    print(f"    - Validation: {info['validation_size']:,} cặp câu")
+    print(f"    - Test: {info['test_size']:,} cặp câu")
+    print(f"    - Total: {info['total_size']:,} cặp câu")
+
+def main():
+    """
+    Main function: Xử lý dữ liệu Hospital
+    """
+    print("\n" + "="*70)
+    print("PREPARE HOSPITAL DATA FOR FINETUNING")
+    print("="*70)
+    
+    # Kiểm tra thư mục Hospital
+    if not HOSPITAL_DIR.exists():
+        print(f"❌ Không tìm thấy thư mục: {HOSPITAL_DIR}")
+        return
+    
+    # Tạo thư mục processed nếu chưa có
+    PROCESSED_DATA_DIR.mkdir(parents=True, exist_ok=True)
+    
+    # Load dữ liệu
+    hospital_data = load_hospital_data()
+    
+    # Encode dữ liệu
+    encoded_data, vi_vocab, en_vocab = encode_hospital_data(hospital_data)
+    
+    # Lưu dữ liệu
+    save_hospital_data(encoded_data)
+    
+    print("\n" + "="*70)
+    print("✓ HOÀN TẤT XỬ LÝ DỮ LIỆU HOSPITAL")
+    print("="*70)
+    print(f"\nDữ liệu đã sẵn sàng để finetune!")
+    print(f"  - Encoded data: {PROCESSED_DATA_DIR / 'hospital_data_encoded.pkl'}")
+    print(f"  - Info: {PROCESSED_DATA_DIR / 'hospital_data_info.json'}")
+
+if __name__ == "__main__":
+    main()
+

src/shared_vocab_utils.py

@@ -0,0 +1,338 @@
+"""
+UTILITY FUNCTIONS CHO SHARED VOCABULARY
+Helper functions để load và sử dụng shared tokenizer
+"""
+
+import json
+import pickle
+import re
+from pathlib import Path
+from tokenizers import Tokenizer
+from typing import Tuple, Optional
+
+# ============================================================================
+# PATH CONFIGURATION
+# ============================================================================
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+PROCESSED_DATA_DIR = PROJECT_ROOT / 'data' / 'processed'
+
+# ============================================================================
+# LOAD SHARED VOCABULARY
+# ============================================================================
+
+def load_shared_tokenizer(tokenizer_path: Optional[Path] = None) -> Tokenizer:
+    """
+    Load shared tokenizer từ file
+    
+    Args:
+        tokenizer_path: Đường dẫn đến tokenizer_shared.json (mặc định tự tìm)
+        
+    Returns:
+        tokenizer: Loaded tokenizer
+    """
+    if tokenizer_path is None:
+        tokenizer_path = PROCESSED_DATA_DIR / 'tokenizer_shared.json'
+    
+    if not tokenizer_path.exists():
+        raise FileNotFoundError(
+            f"Không tìm thấy tokenizer_shared.json tại {tokenizer_path}!\n"
+            f"Vui lòng chạy: python src/1_build_shared_vocab.py trước"
+        )
+    
+    tokenizer = Tokenizer.from_file(str(tokenizer_path))
+    return tokenizer
+
+def load_shared_vocab_info(info_path: Optional[Path] = None) -> dict:
+    """
+    Load thông tin về shared vocabulary
+    
+    Args:
+        info_path: Đường dẫn đến shared_vocab_info.json (mặc định tự tìm)
+        
+    Returns:
+        info: Dictionary chứa vocab_size, sos_id, eos_id, pad_id, unk_id
+    """
+    if info_path is None:
+        info_path = PROCESSED_DATA_DIR / 'shared_vocab_info.json'
+    
+    if not info_path.exists():
+        # Tạo từ tokenizer nếu chưa có
+        tokenizer = load_shared_tokenizer()
+        info = {
+            'vocab_size': tokenizer.get_vocab_size(),
+            'sos_id': tokenizer.token_to_id("<sos>"),
+            'eos_id': tokenizer.token_to_id("<eos>"),
+            'pad_id': tokenizer.token_to_id("<pad>"),
+            'unk_id': tokenizer.token_to_id("<unk>"),
+            'tokenizer_path': str(PROCESSED_DATA_DIR / 'tokenizer_shared.json')
+        }
+        return info
+    
+    with open(info_path, 'r', encoding='utf-8') as f:
+        info = json.load(f)
+    
+    return info
+
+def load_shared_processed_data(data_path: Optional[Path] = None, use_bidirectional: bool = False) -> dict:
+    """
+    Load processed data đã encode với shared vocabulary
+    
+    Args:
+        data_path: Đường dẫn đến processed_data file (mặc định tự tìm)
+        use_bidirectional: Nếu True, load bidirectional dataset (có cả vi→en và en→vi)
+        
+    Returns:
+        processed_data: Dict với keys 'train', 'validation', 'test'
+                       Mỗi value là list of (src_ids, tgt_ids) tuples
+    """
+    if data_path is None:
+        if use_bidirectional:
+            data_path = PROCESSED_DATA_DIR / 'processed_data_bidirectional.pkl'
+            if not data_path.exists():
+                # Fallback to regular dataset if bidirectional doesn't exist
+                print("⚠️  Bidirectional dataset không tồn tại, dùng dataset thường")
+                data_path = PROCESSED_DATA_DIR / 'processed_data_shared.pkl'
+        else:
+            data_path = PROCESSED_DATA_DIR / 'processed_data_shared.pkl'
+    
+    if not data_path.exists():
+        raise FileNotFoundError(
+            f"Không tìm thấy data file tại {data_path}!\n"
+            f"Vui lòng chạy: python src/2_encode_data.py trước"
+        )
+    
+    with open(data_path, 'rb') as f:
+        processed_data = pickle.load(f)
+    
+    if use_bidirectional and 'bidirectional' in str(data_path):
+        print(f"✓ Loaded bidirectional dataset từ {data_path.name}")
+    else:
+        print(f"✓ Loaded dataset từ {data_path.name}")
+    
+    return processed_data
+
+# ============================================================================
+# POST-PROCESSING: CLEAN DECODED OUTPUT
+# ============================================================================
+
+def clean_decoded_output(text: str) -> str:
+    """
+    Clean decoded output để loại bỏ padding artifacts và các lỗi kỹ thuật
+    
+    Fixes:
+    1. Loại bỏ khoảng trắng thừa
+    2. Loại bỏ ký tự lạ (như ‹¶, <pad>, etc.)
+    3. Normalize spacing quanh dấu câu
+    4. Strip leading/trailing whitespace
+    
+    Args:
+        text: Raw decoded text
+        
+    Returns:
+        cleaned: Cleaned text
+    """
+    if not text:
+        return ""
+    
+    # Loại bỏ các ký tự đặc biệt không mong muốn
+    # Loại bỏ các ký tự control và non-printable (trừ space, newline, tab)
+    text = re.sub(r'[\x00-\x08\x0B-\x0C\x0E-\x1F\x7F-\x9F]', '', text)
+    
+    # Loại bỏ các ký tự lạ như ‹¶ và các ký tự Unicode không hợp lệ
+    text = re.sub(r'[ܦ]', '', text)
+    
+    # Loại bỏ các token đặc biệt nếu còn sót
+    text = text.replace('<pad>', '').replace('</s>', '').replace('<s>', '')
+    text = text.replace('<eos>', '').replace('<sos>', '').replace('<unk>', '')
+    
+    # Normalize spacing quanh dấu câu (loại bỏ space trước dấu câu)
+    text = re.sub(r'\s+([,.!?;:])', r'\1', text)
+    # Thêm space sau dấu câu nếu chưa có (trừ khi là cuối câu)
+    text = re.sub(r'([,.!?;:])([^\s])', r'\1 \2', text)
+    
+    # Loại bỏ multiple spaces
+    text = re.sub(r' +', ' ', text)
+    
+    # Loại bỏ space ở đầu và cuối
+    text = text.strip()
+    
+    # Loại bỏ space thừa ở đầu câu (nếu có)
+    text = text.lstrip()
+    
+    return text
+
+# ============================================================================
+# ENCODE/DECODE HELPERS
+# ============================================================================
+
+class SharedVocabulary:
+    """
+    Wrapper class để dùng shared tokenizer như Vocabulary cũ
+    Tương thích với code hiện tại
+    """
+    def __init__(self, tokenizer: Tokenizer):
+        self.tokenizer = tokenizer
+        self.PAD_TOKEN = "<pad>"
+        self.SOS_TOKEN = "<sos>"
+        self.EOS_TOKEN = "<eos>"
+        self.UNK_TOKEN = "<unk>"
+        
+        self.PAD_IDX = tokenizer.token_to_id(self.PAD_TOKEN) or 0
+        self.SOS_IDX = tokenizer.token_to_id(self.SOS_TOKEN) or 1
+        self.EOS_IDX = tokenizer.token_to_id(self.EOS_TOKEN) or 2
+        self.UNK_IDX = tokenizer.token_to_id(self.UNK_TOKEN) or 3
+    
+    def encode(self, sentence: str) -> list:
+        """
+        Encode câu thành list of token IDs
+        
+        Args:
+            sentence: Input sentence (string)
+            
+        Returns:
+            tokens: List of token IDs [SOS_ID, ...token_ids..., EOS_ID]
+        """
+        # Clean sentence (lowercase, normalize spaces)
+        sentence = sentence.lower().strip()
+        
+        # Encode với tokenizer
+        encoded = self.tokenizer.encode(sentence)
+        token_ids = encoded.ids
+        
+        # Thêm SOS và EOS
+        return [self.SOS_IDX] + token_ids + [self.EOS_IDX]
+    
+    def decode(self, token_ids: list) -> str:
+        """
+        Decode list of token IDs thành câu
+        
+        Args:
+            token_ids: List of token IDs
+            
+        Returns:
+            sentence: Decoded sentence (string)
+        """
+        # Loại bỏ SOS, EOS, PAD
+        filtered_ids = [
+            idx for idx in token_ids
+            if idx not in [self.PAD_IDX, self.SOS_IDX, self.EOS_IDX]
+        ]
+        
+        if not filtered_ids:
+            return ""
+        
+        # Decode với tokenizer
+        decoded = self.tokenizer.decode(filtered_ids, skip_special_tokens=True)
+        
+        # Post-processing: Clean output
+        decoded = clean_decoded_output(decoded)
+        
+        return decoded
+    
+    def __len__(self):
+        """Kích thước vocabulary"""
+        return self.tokenizer.get_vocab_size()
+    
+    def token_to_id(self, token: str) -> int:
+        """Convert token string to ID"""
+        return self.tokenizer.token_to_id(token) or self.UNK_IDX
+    
+    def id_to_token(self, idx: int) -> str:
+        """Convert ID to token string"""
+        return self.tokenizer.id_to_token(idx) or self.UNK_TOKEN
+
+def create_shared_vocab_wrapper() -> Tuple[SharedVocabulary, SharedVocabulary]:
+    """
+    Tạo wrapper cho shared vocabulary (tương thích với code cũ)
+    
+    Returns:
+        vi_vocab, en_vocab: Cả 2 đều là SharedVocabulary (cùng tokenizer)
+    """
+    tokenizer = load_shared_tokenizer()
+    
+    # Cả 2 ngôn ngữ dùng chung tokenizer
+    vi_vocab = SharedVocabulary(tokenizer)
+    en_vocab = SharedVocabulary(tokenizer)
+    
+    return vi_vocab, en_vocab
+
+# ============================================================================
+# CHECK SHARED VOCAB SETUP
+# ============================================================================
+
+def check_shared_vocab_setup() -> bool:
+    """
+    Kiểm tra xem shared vocabulary đã được setup chưa
+    
+    Returns:
+        is_setup: True nếu đã setup đầy đủ
+    """
+    tokenizer_path = PROCESSED_DATA_DIR / 'tokenizer_shared.json'
+    data_path = PROCESSED_DATA_DIR / 'processed_data_shared.pkl'
+    
+    return tokenizer_path.exists() and data_path.exists()
+
+def print_shared_vocab_status():
+    """
+    In trạng thái setup của shared vocabulary
+    """
+    print("="*70)
+    print("KIỂM TRA SHARED VOCABULARY SETUP")
+    print("="*70)
+    
+    tokenizer_path = PROCESSED_DATA_DIR / 'tokenizer_shared.json'
+    data_path = PROCESSED_DATA_DIR / 'processed_data_shared.pkl'
+    info_path = PROCESSED_DATA_DIR / 'shared_vocab_info.json'
+    
+    print(f"\n1. Tokenizer:")
+    if tokenizer_path.exists():
+        print(f"   ✓ {tokenizer_path}")
+        try:
+            tokenizer = load_shared_tokenizer()
+            print(f"   ✓ Vocab size: {tokenizer.get_vocab_size()}")
+        except Exception as e:
+            print(f"   ✗ Lỗi khi load: {e}")
+    else:
+        print(f"   ✗ Không tìm thấy: {tokenizer_path}")
+        print(f"   → Chạy: python src/1_build_shared_vocab.py")
+    
+    print(f"\n2. Processed Data:")
+    if data_path.exists():
+        print(f"   ✓ {data_path}")
+        try:
+            data = load_shared_processed_data()
+            print(f"   ✓ Train: {len(data.get('train', []))} cặp câu")
+            print(f"   ✓ Validation: {len(data.get('validation', []))} cặp câu")
+            print(f"   ✓ Test: {len(data.get('test', []))} cặp câu")
+        except Exception as e:
+            print(f"   ✗ Lỗi khi load: {e}")
+    else:
+        print(f"   ✗ Không tìm thấy: {data_path}")
+        print(f"   → Chạy: python src/2_encode_data.py")
+    
+    print(f"\n3. Info File:")
+    if info_path.exists():
+        print(f"   ✓ {info_path}")
+    else:
+        print(f"   ⚠️  Chưa có (sẽ tự tạo khi cần)")
+    
+    print("\n" + "="*70)
+    
+    if check_shared_vocab_setup():
+        print("✓ SHARED VOCABULARY ĐÃ SẴN SÀNG!")
+    else:
+        print("⚠️  SHARED VOCABULARY CHƯA ĐƯỢC SETUP")
+        print("\n📝 Các bước cần làm:")
+        print("   1. python src/1_build_shared_vocab.py")
+        print("   2. python src/2_encode_data.py")
+    
+    print("="*70)
+
+# ============================================================================
+# MAIN
+# ============================================================================
+
+if __name__ == "__main__":
+    print_shared_vocab_status()
+

src/train_mtet_bidirectional.py

@@ -0,0 +1,347 @@
+"""
+Train model với mtet_bidirectional dataset
+"""
+
+import os
+# Set CUDA memory allocation để tránh fragmentation
+os.environ['PYTORCH_ALLOC_CONF'] = 'expandable_segments:True'
+
+import torch
+from pathlib import Path
+import sys
+import argparse
+
+# Add parent directory to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from src.complete_transformer import TransformerShared, get_model_config
+from src.shared_vocab_utils import load_shared_tokenizer, load_shared_vocab_info
+from src.dataloader_module import create_dataloaders_with_bucketing
+from src.training_module import train_model
+from src.gpu_safety import get_gpu_memory_info, check_and_adjust_batch_size, clear_gpu_cache
+import pickle
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+PROCESSED_DATA_DIR = PROJECT_ROOT / 'data' / 'processed'
+CHECKPOINT_DIR = PROJECT_ROOT / 'checkpoints'
+
+def load_mtet_bidirectional_data():
+    """Load processed mtet bidirectional data"""
+    # Dùng bản dữ liệu đã được clean & bidirectional từ mtet_cleaned.csv
+    data_path = PROCESSED_DATA_DIR / "processed_data_mtet_bidirectional_cleaned.pkl"
+    
+    if not data_path.exists():
+        raise FileNotFoundError(
+            f"Không tìm thấy processed data tại {data_path}!\n"
+            f"Vui lòng chạy: python src/encode_mtet_bidirectional.py trước"
+        )
+    
+    print(f"📂 Loading data từ: {data_path}")
+    with open(data_path, 'rb') as f:
+        data = pickle.load(f)
+    
+    # Convert format từ list of tuples sang dict với 'src' và 'tgt'
+    converted_data = {}
+    for split in ['train', 'validation', 'test']:
+        if split in data:
+            src_list = [item[0] for item in data[split]]
+            tgt_list = [item[1] for item in data[split]]
+            converted_data[split] = {
+                'src': src_list,
+                'tgt': tgt_list
+            }
+        else:
+            converted_data[split] = {'src': [], 'tgt': []}
+    
+    return converted_data
+
+def main():
+    parser = argparse.ArgumentParser(description="Train với mtet_bidirectional dataset")
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        default=30,
+        help="Số epochs (mặc định: 30)",
+    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=32,
+        help="Batch size (mặc định: 32, giảm để tránh OOM)",
+    )
+    parser.add_argument(
+        "--model_size",
+        type=str,
+        default="custom_25m",
+        choices=["custom_25m", "base", "small"],
+        help="Kích thước model (mặc định: custom_25m)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda",
+        help="Device (cuda/cpu, mặc định: cuda)",
+    )
+    parser.add_argument(
+        "--precision",
+        type=str,
+        default="bf16",
+        choices=["fp32", "amp", "bf16"],
+        help="Precision (mặc định: bf16)",
+    )
+    parser.add_argument(
+        "--grad_accum_steps",
+        type=int,
+        default=1,
+        help="Gradient accumulation steps (mặc định: 1)",
+    )
+    parser.add_argument(
+        "--warmup_steps",
+        type=int,
+        default=4000,
+        help="Warmup steps (mặc định: 4000 steps)",
+    )
+    parser.add_argument(
+        "--target_max_lr",
+        type=float,
+        default=3e-4,
+        help="Đỉnh learning rate mong muốn cho scheduler (mặc định: 3e-4)",
+    )
+    parser.add_argument(
+        "--label_smoothing",
+        type=float,
+        default=0.1,
+        help="Label smoothing (mặc định: 0.1)",
+    )
+    parser.add_argument(
+        "--checkpoint_dir",
+        type=str,
+        default=None,
+        help="Checkpoint directory (mặc định: checkpoints)",
+    )
+    parser.add_argument(
+        "--resume_from",
+        type=str,
+        default=None,
+        help="Resume từ checkpoint (đường dẫn đến file .pt)",
+    )
+    parser.add_argument(
+        "--no_compile",
+        action="store_true",
+        help="Tắt torch.compile (mặc định: BẬT để tăng tốc)",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=None,
+        help="Số workers cho DataLoader (mặc định: 2 cho CUDA, 0 cho CPU)",
+    )
+    
+    args = parser.parse_args()
+    
+    # Device
+    device = torch.device(args.device)
+    if device.type == 'cuda' and not torch.cuda.is_available():
+        print("⚠️  CUDA không khả dụng, dùng CPU")
+        device = torch.device('cpu')
+    
+    print("="*70)
+    print("TRAIN MODEL VỚI MTET BIDIRECTIONAL DATASET")
+    print("="*70)
+    print(f"Device: {device}")
+    print(f"Model size: {args.model_size}")
+    
+    # Kiểm tra GPU memory và điều chỉnh batch size nếu cần
+    if device.type == 'cuda':
+        clear_gpu_cache()  # Xóa cache trước khi kiểm tra
+        mem_info = get_gpu_memory_info()
+        if mem_info:
+            print(f"\n📊 GPU Memory Info:")
+            print(f"   Total: {mem_info['total']:.2f} GB")
+            print(f"   Free: {mem_info['free']:.2f} GB")
+            print(f"   Usage: {mem_info['usage_percent']:.1f}%")
+        
+        # Kiểm tra và điều chỉnh batch size
+        safe_batch_size, warning = check_and_adjust_batch_size(args.batch_size)
+        if warning:
+            print(f"\n{warning}")
+            response = input(f"\n❓ Có muốn giảm batch_size xuống {safe_batch_size} không? (y/n, mặc định: n): ").strip().lower()
+            if response == 'y':
+                args.batch_size = safe_batch_size
+                print(f"✓ Đã giảm batch_size xuống {safe_batch_size}")
+            else:
+                print(f"⚠️  Giữ nguyên batch_size={args.batch_size}, có thể gặp OOM")
+        else:
+            print(f"✓ Batch size {args.batch_size} an toàn với GPU memory hiện tại")
+    
+    print(f"\nBatch size: {args.batch_size}")
+    print(f"Gradient accumulation: {args.grad_accum_steps}")
+    effective_batch = args.batch_size * args.grad_accum_steps
+    print(f"Effective batch size: {effective_batch} ({args.batch_size} × {args.grad_accum_steps})")
+    print(f"Epochs: {args.num_epochs}")
+    print(f"Precision: {args.precision}")
+    print()
+    print("🔄 BIDIRECTIONAL TRAINING:")
+    print("   Model sẽ học ĐỒNG THỜI cả vi→en và en→vi")
+    print("   Dataset đã hỗn hợp cả hai chiều (~50% mỗi chiều)")
+    print("="*70 + "\n")
+    
+    # Load vocab info
+    print("📚 Loading vocabulary...")
+    vocab_info = load_shared_vocab_info()
+    vocab_size = vocab_info['vocab_size']
+    pad_idx = vocab_info['pad_id']
+    print(f"✓ Vocab size: {vocab_size:,}")
+    print(f"✓ Pad ID: {pad_idx}\n")
+    
+    # Load data
+    print("📂 Loading data...")
+    processed_data = load_mtet_bidirectional_data()
+    
+    print(f"✓ Train: {len(processed_data['train']['src']):,} cặp câu")
+    print(f"✓ Validation: {len(processed_data['validation']['src']):,} cặp câu")
+    print(f"✓ Test: {len(processed_data['test']['src']):,} cặp câu\n")
+    
+    # Create dataloaders
+    print("🔄 Creating dataloaders...")
+    # Tối ưu num_workers: 8 quá nhiều gây overhead, dùng 2-4 thay vì
+    # Với dataset lớn (7M samples), num_workers cao gây lag do context switching
+    if args.num_workers is None:
+        num_workers = 2 if device.type == 'cuda' else 0  # Mặc định: 2 cho CUDA
+    else:
+        num_workers = args.num_workers
+    train_loader, val_loader, test_loader = create_dataloaders_with_bucketing(
+        processed_data,
+        batch_size=args.batch_size,
+        num_workers=num_workers,
+    )
+    print(f"✓ Dataloaders created với BucketSampler (num_workers={num_workers})\n")
+    
+    # Create model
+    print("🔨 Creating model...")
+    model_config = get_model_config(args.model_size)
+    
+    model = TransformerShared(
+        vocab_size=vocab_size,
+        d_model=model_config['d_model'],
+        n_layers=model_config['n_layers'],
+        n_heads=model_config['n_heads'],
+        d_ff=model_config['d_ff'],
+        dropout=model_config['dropout'],
+        pad_idx=pad_idx,
+        use_weight_tying=True
+    )
+    
+    model = model.to(device)
+    
+    # Tối ưu: Compile model với torch.compile (PyTorch 2.0+)
+    # LƯU Ý: torch.compile KHÔNG tương thích với gradient accumulation
+    # Nếu grad_accum_steps > 1, tự động tắt compile để tránh lỗi CUDAGraphs
+    compile_model = not args.no_compile and args.grad_accum_steps == 1
+    if compile_model and hasattr(torch, 'compile') and device.type == 'cuda':
+        print("🔧 Compiling model với torch.compile để tăng tốc...")
+        print("   ⚠️  Lưu ý: Có thể chậm ở đầu training (~10-20% epoch đầu)")
+        print("   💡 Tăng tốc ~20-30% sau khi compile xong")
+        try:
+            # Dùng mode='default' thay vì 'reduce-overhead' để giảm memory
+            model = torch.compile(model, mode='default')
+            print("✓ Model đã được compile (mode='default' để giảm memory)\n")
+        except Exception as e:
+            print(f"⚠️  Không thể compile model: {e}")
+            print("   Tiếp tục với model không compile\n")
+    else:
+        if args.grad_accum_steps > 1:
+            print("⚠️  torch.compile đã tắt (không tương thích với gradient accumulation)\n")
+        elif args.no_compile:
+            print("⚠️  torch.compile đã tắt (--no_compile)\n")
+        else:
+            print("⚠️  torch.compile không khả dụng\n")
+    
+    # Count parameters
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    model_size_mb = total_params * 4 / (1024**2)  # float32
+    print(f"✓ Model created")
+    print(f"  - Total parameters: {total_params:,}")
+    print(f"  - Trainable parameters: {trainable_params:,}")
+    print(f"  - Model size: ~{model_size_mb:.2f} MB (float32)")
+    
+    # Kiểm tra lại GPU memory sau khi load model
+    if device.type == 'cuda':
+        clear_gpu_cache()
+        mem_info = get_gpu_memory_info()
+        if mem_info:
+            print(f"\n📊 GPU Memory sau khi load model:")
+            print(f"   Allocated: {mem_info['allocated']:.2f} GB")
+            print(f"   Reserved: {mem_info['reserved']:.2f} GB")
+            print(f"   Free: {mem_info['free']:.2f} GB")
+            print(f"   Usage: {mem_info['usage_percent']:.1f}%")
+            
+            # Kiểm tra lại batch size sau khi load model
+            safe_batch_size, warning = check_and_adjust_batch_size(args.batch_size)
+            if warning and safe_batch_size < args.batch_size:
+                print(f"\n{warning}")
+                print(f"⚠️  Khuyến nghị: Giảm batch_size xuống {safe_batch_size} để tránh OOM")
+                print(f"   Hoặc tiếp tục với batch_size={args.batch_size} (có thể gặp OOM)")
+    
+    print()
+    
+    # Checkpoint directory
+    checkpoint_dir = Path(args.checkpoint_dir) if args.checkpoint_dir else CHECKPOINT_DIR
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Resume từ checkpoint nếu có
+    resume_from = None
+    if args.resume_from:
+        resume_path = Path(args.resume_from)
+        if resume_path.exists():
+            resume_from = str(resume_path)
+            print(f"🔄 Sẽ resume từ: {resume_from}")
+        else:
+            print(f"⚠️  Không tìm thấy checkpoint: {resume_from}")
+            print("   Sẽ train từ đầu")
+    else:
+        # Tự động tìm checkpoint mới nhất nếu có
+        checkpoints = list(checkpoint_dir.glob('checkpoint_epoch_*.pt'))
+        if checkpoints:
+            latest = max(checkpoints, key=lambda p: int(p.stem.split('_')[-1]))
+            print(f"💡 Tìm thấy checkpoint cũ: {latest.name}")
+            print(f"   Để resume, dùng: --resume_from {latest}")
+            print(f"   Hoặc xóa checkpoints cũ để train từ đầu\n")
+    
+    # Clear cache trước khi training
+    if device.type == 'cuda':
+        clear_gpu_cache()
+        print("✓ Đã clear GPU cache trước khi training\n")
+    
+    # Train
+    print("🚀 Starting training...")
+    print("="*70 + "\n")
+    
+    history = train_model(
+        model=model,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        num_epochs=args.num_epochs,
+        device=device,
+        d_model=model_config["d_model"],
+        warmup_steps=args.warmup_steps,
+        label_smoothing=args.label_smoothing,
+        grad_accum_steps=args.grad_accum_steps,
+        precision=args.precision,
+        checkpoint_dir=checkpoint_dir,
+        save_every=1,
+        resume_from=resume_from,
+        target_max_lr=args.target_max_lr,
+    )
+    
+    print("\n" + "="*70)
+    print("✅ TRAINING HOÀN TẤT!")
+    print("="*70)
+    print(f"Checkpoints được lưu tại: {checkpoint_dir}")
+    print(f"Best model: {checkpoint_dir / 'best_model.pt'}")
+    print("="*70)
+
+if __name__ == '__main__':
+    main()
+

src/training_module.py

@@ -0,0 +1,771 @@
+from __future__ import annotations
+
+"""
+PHẦN C: HUẤN LUYỆN VÀ ĐÁNH GIÁ (FIXED)
+Training loop, Loss function, Optimizer, Learning Rate Scheduler
+"""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.optim.lr_scheduler import _LRScheduler
+from torch.cuda.amp import GradScaler
+import math
+import time
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+from pathlib import Path
+
+# ============================================================================
+# 1. LABEL SMOOTHING CROSS ENTROPY LOSS (FIXED)
+# ============================================================================
+
+class LabelSmoothingLoss(nn.Module):
+    """
+    Label Smoothing Cross Entropy Loss
+    
+    Làm mịn nhãn để tránh overfitting:
+    - True label: 1 - smoothing
+    - Other labels: smoothing / (vocab_size - 1)
+    
+    Args:
+        vocab_size: Kích thước vocabulary
+        pad_idx: Index của padding token (không tính loss)
+        smoothing: Label smoothing factor (mặc định 0.1)
+    """
+    def __init__(self, vocab_size, pad_idx=0, smoothing=0.1):
+        super().__init__()
+        
+        self.vocab_size = vocab_size
+        self.pad_idx = pad_idx
+        self.smoothing = smoothing
+        self.confidence = 1.0 - smoothing
+        
+    def forward(self, pred, target):
+        """
+        Args:
+            pred: Predictions [batch_size, seq_len, vocab_size]
+            target: Target labels [batch_size, seq_len]
+            
+        Returns:
+            loss: Scalar loss value
+        """
+        batch_size, seq_len, vocab_size = pred.size()
+        
+        # Reshape - FIXED: use .reshape() instead of .view() for non-contiguous tensors
+        pred = pred.reshape(-1, vocab_size)  # [batch*seq_len, vocab_size]
+        target = target.reshape(-1)  # [batch*seq_len]
+        
+        # Log softmax
+        log_probs = torch.log_softmax(pred, dim=-1)
+        
+        # Tạo smoothed labels
+        smoothed_targets = torch.zeros_like(log_probs)
+        smoothed_targets.fill_(self.smoothing / (vocab_size - 2))  # -2 vì trừ true label và pad
+        smoothed_targets.scatter_(1, target.unsqueeze(1), self.confidence)
+        smoothed_targets[:, self.pad_idx] = 0  # Không tính loss cho padding
+        
+        # Mask padding
+        mask = (target != self.pad_idx).float()
+        
+        # Tính loss
+        loss = -(smoothed_targets * log_probs).sum(dim=-1)
+        loss = (loss * mask).sum() / mask.sum()
+        
+        return loss
+
+# ============================================================================
+# 2. LEARNING RATE SCHEDULER - WARMUP + DECAY
+# ============================================================================
+
+class TransformerLRScheduler(_LRScheduler):
+    """
+    Learning Rate Scheduler với Warmup + Decay cho Adam optimizer
+    
+    lr = d_model^(-0.5) * min(step^(-0.5), step * warmup_steps^(-1.5)) * decay_factor
+    
+    Args:
+        optimizer: PyTorch optimizer
+        d_model: Dimension của model
+        warmup_steps: Số steps warmup (mặc định 4000)
+        factor: Scaling factor (điều chỉnh để đạt max LR mong muốn)
+        decay_factor: Decay factor sau warmup (mặc định 0.8, giảm 20% mỗi epoch)
+    """
+    def __init__(self, optimizer, d_model, warmup_steps=10000, factor=0.98, decay_factor=0.98):
+        self.d_model = d_model
+        self.warmup_steps = warmup_steps
+        # factor chỉnh để peak LR ≈ 0.0005 (5e-4)
+        # Tính toán: factor = 0.0005 / (384^(-0.5) * 10000^(-0.5)) ≈ 0.98
+        self.factor = factor
+        self.decay_factor = decay_factor  # Giảm từ 0.95 → 0.98 (decay chậm hơn)
+        self.num_steps = 0
+        self.last_epoch = 0
+        
+        super().__init__(optimizer)
+        
+    def get_lr(self):
+        """
+        Tính learning rate cho step hiện tại với decay
+        """
+        self.num_steps += 1
+        
+        # Base Transformer schedule
+        base_lr = self.factor * (
+            self.d_model ** (-0.5) *
+            min(self.num_steps ** (-0.5), 
+                self.num_steps * self.warmup_steps ** (-1.5))
+        )
+        
+        # Apply decay after warmup (giảm dần sau khi warmup xong)
+        # Decay interval: warmup_steps (thay vì warmup_steps//2) để decay chậm hơn
+        if self.num_steps > self.warmup_steps:
+            decay_steps = (self.num_steps - self.warmup_steps) // self.warmup_steps
+            lr = base_lr * (self.decay_factor ** decay_steps)
+        else:
+            lr = base_lr
+        
+        return [lr for _ in self.base_lrs]
+
+# ============================================================================
+# 3. PERPLEXITY METRIC
+# ============================================================================
+
+def calculate_perplexity(loss):
+    """
+    Tính Perplexity từ loss
+    
+    Perplexity = exp(loss)
+    
+    Args:
+        loss: Cross-entropy loss
+        
+    Returns:
+        perplexity: Perplexity score
+    """
+    return math.exp(min(loss, 100))  # Cap để tránh overflow
+
+# ============================================================================
+# 4. TRAINING FUNCTION
+# ============================================================================
+
+def train_epoch(
+    model,
+    train_loader,
+    optimizer,
+    scheduler,
+    criterion,
+    device,
+    epoch,
+    grad_accum_steps=1,
+    use_amp=False,
+    use_bf16=False,
+    scaler: GradScaler | None = None,
+):
+    """
+    Train một epoch
+    
+    Args:
+        model: Transformer model
+        train_loader: Training DataLoader
+        optimizer: Optimizer
+        scheduler: Learning rate scheduler
+        criterion: Loss function
+        device: Device (cuda/cpu)
+        epoch: Epoch number
+        use_amp: Use float16 mixed precision (requires GradScaler)
+        use_bf16: Use bfloat16 mixed precision (no GradScaler needed)
+        scaler: GradScaler for float16 (not needed for bf16)
+        
+    Returns:
+        avg_loss: Average loss
+        avg_perplexity: Average perplexity
+    """
+    model.train()
+    
+    total_loss = 0
+    total_tokens = 0
+    
+    progress_bar = tqdm(train_loader, desc=f'Epoch {epoch}', ncols=120, file=None, dynamic_ncols=False)
+    
+    optimizer.zero_grad(set_to_none=True)
+    # Chỉ dùng GradScaler cho float16, không cần cho bf16
+    if use_amp and not use_bf16:
+        scaler = scaler or GradScaler(enabled=True)
+    num_batches = len(train_loader)
+
+    # Import GPU safety functions
+    try:
+        from src.gpu_safety import clear_gpu_cache, check_memory_spike
+    except ImportError:
+        def clear_gpu_cache():
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        def check_memory_spike(threshold=0.9):
+            return False, None
+    
+    skipped_batches = 0
+    skipped_long_seq = 0  # Đếm riêng skip do sequence quá dài
+    skipped_oom = 0  # Đếm riêng skip do OOM
+    max_skipped_batches = 20  # Giới hạn tổng số batch skip
+    max_skipped_long_seq = 15  # Giới hạn skip do sequence dài
+    processed_batches = 0
+
+    for batch_idx, (src, tgt, src_len, tgt_len) in enumerate(progress_bar):
+        # Kiểm tra giới hạn skip để tránh vòng lặp vô tận
+        if skipped_batches >= max_skipped_batches:
+            print(f"\n⚠️  Đã skip {skipped_batches} batches (giới hạn: {max_skipped_batches})")
+            print(f"   - Skip do sequence dài: {skipped_long_seq}")
+            print(f"   - Skip do OOM: {skipped_oom}")
+            print("   Dừng training để tránh vòng lặp vô tận")
+            break
+        
+        # Kiểm tra nếu skip quá nhiều do sequence dài
+        if skipped_long_seq >= max_skipped_long_seq:
+            print(f"\n⚠️  Đã skip {skipped_long_seq} batches do sequence quá dài (giới hạn: {max_skipped_long_seq})")
+            print("   Khuyến nghị: Kiểm tra lại dataset hoặc tăng max_seq_len")
+            break
+        
+        try:
+            # Kiểm tra memory spike trước khi xử lý batch
+            if device.type == 'cuda' and batch_idx % 100 == 0:
+                is_spike, mem_info = check_memory_spike(threshold=0.9)
+                if is_spike and mem_info:
+                    print(f"\n⚠️  GPU Memory spike detected: {mem_info['usage_percent']:.1f}%")
+                    print(f"   Free: {mem_info['free']:.2f}GB")
+                    clear_gpu_cache()
+            
+            # Kiểm tra batch size quá lớn (sequence dài)
+            max_seq_len = max(src.size(1), tgt.size(1))
+            if max_seq_len > 200:  # Batch có sequence quá dài
+                skipped_batches += 1
+                skipped_long_seq += 1
+                if skipped_long_seq % 5 == 0:  # Log mỗi 5 batches
+                    print(f"\n⚠️  Đã skip {skipped_long_seq} batches do sequence quá dài (>200 tokens)")
+                continue
+            
+            # Move to device
+            src = src.to(device)
+            tgt = tgt.to(device)
+            
+            # Target input (bỏ token cuối) và output (bỏ token đầu)
+            tgt_input = tgt[:, :-1]
+            tgt_output = tgt[:, 1:]
+            
+            # Autocast với dtype phù hợp (API mới torch.amp.autocast)
+            if use_bf16:
+                with torch.amp.autocast('cuda', enabled=True, dtype=torch.bfloat16):
+                    output = model(src, tgt_input)
+                    loss = criterion(output, tgt_output)
+                    loss_to_backward = loss / grad_accum_steps
+            elif use_amp:
+                with torch.amp.autocast('cuda', enabled=True):
+                    output = model(src, tgt_input)
+                    loss = criterion(output, tgt_output)
+                    loss_to_backward = loss / grad_accum_steps
+            else:
+                output = model(src, tgt_input)
+                loss = criterion(output, tgt_output)
+                loss_to_backward = loss / grad_accum_steps
+            
+            # Backward pass
+            if use_amp and not use_bf16:
+                scaler.scale(loss_to_backward).backward()
+            else:
+                loss_to_backward.backward()
+            
+            should_step = (
+                ((batch_idx + 1) % grad_accum_steps == 0) or (batch_idx + 1 == num_batches)
+            )
+            
+            if should_step:
+                if use_amp and not use_bf16:
+                    scaler.unscale_(optimizer)
+                # Gradient clipping mạnh tay hơn để giảm loss spikes
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)
+                
+                if use_amp and not use_bf16:
+                    scaler.step(optimizer)
+                    scaler.update()
+                else:
+                    optimizer.step()
+                
+                scheduler.step()
+                optimizer.zero_grad(set_to_none=True)
+            
+            # Statistics
+            num_tokens = (tgt_output != criterion.pad_idx).sum().item()
+            total_loss += loss.item() * num_tokens
+            total_tokens += num_tokens
+            processed_batches += 1
+            
+            # Update progress bar
+            current_loss = loss.item()
+            current_ppl = calculate_perplexity(current_loss)
+            current_lr = scheduler.get_last_lr()[0]
+            
+            progress_bar.set_postfix({
+                'loss': f'{current_loss:.4f}',
+                'ppl': f'{current_ppl:.2f}',
+                'lr': f'{current_lr:.6f}',
+                'skip': f'{skipped_batches}'
+            })
+    
+        except RuntimeError as e:
+            # Xử lý OOM error
+            error_msg = str(e)
+            if 'out of memory' in error_msg.lower() or 'cuda' in error_msg.lower():
+                print(f"\n❌ OOM Error tại batch {batch_idx}: {error_msg}")
+                print("   Đang xóa cache và skip batch này...")
+                
+                # Xóa cache
+                clear_gpu_cache()
+                
+                # Skip batch này
+                skipped_batches += 1
+                skipped_oom += 1
+                
+                # Reset gradients
+                optimizer.zero_grad(set_to_none=True)
+                
+                # Log mỗi 5 OOM errors
+                if skipped_oom % 5 == 0:
+                    print(f"\n⚠️  Đã skip {skipped_oom} batches do OOM")
+                
+                # Sau khi xử lý OOM, skip batch hiện tại
+                continue
+                continue
+            else:
+                # Lỗi khác, raise lại
+                raise
+    
+    if total_tokens > 0:
+        avg_loss = total_loss / total_tokens
+        avg_perplexity = calculate_perplexity(avg_loss)
+    else:
+        avg_loss = float('inf')
+        avg_perplexity = float('inf')
+        print("\n⚠️  Không có batch nào được xử lý thành công!")
+    
+    # Thống kê chi tiết về skip
+    if skipped_batches > 0:
+        print(f"\n📊 THỐNG KÊ SKIP BATCHES:")
+        print(f"   Tổng số batches đã xử lý: {processed_batches}")
+        print(f"   Tổng số batches đã skip: {skipped_batches}")
+        print(f"     - Skip do sequence quá dài (>200 tokens): {skipped_long_seq}")
+        print(f"     - Skip do OOM: {skipped_oom}")
+        if processed_batches > 0:
+            success_rate = processed_batches / (processed_batches + skipped_batches) * 100
+            print(f"   Tỷ lệ thành công: {success_rate:.1f}%")
+    
+    return avg_loss, avg_perplexity
+
+# ============================================================================
+# 5. VALIDATION FUNCTION
+# ============================================================================
+
+def validate(model, val_loader, criterion, device, use_amp=False, use_bf16=False):
+    """
+    Đánh giá trên validation set
+    
+    Args:
+        model: Transformer model
+        val_loader: Validation DataLoader
+        criterion: Loss function
+        device: Device (cuda/cpu)
+        use_amp: Use float16 mixed precision
+        use_bf16: Use bfloat16 mixed precision
+        
+    Returns:
+        avg_loss: Average loss
+        avg_perplexity: Average perplexity
+    """
+    # Import GPU safety functions
+    try:
+        from src.gpu_safety import clear_gpu_cache, check_memory_spike
+    except ImportError:
+        def clear_gpu_cache():
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        def check_memory_spike(threshold=0.9):
+            return False, None
+    
+    # Clear cache trước khi validation
+    if device.type == 'cuda':
+        clear_gpu_cache()
+    
+    model.eval()
+    
+    total_loss = 0
+    total_tokens = 0
+    skipped_batches = 0
+    max_skipped_batches = 10  # Giới hạn số batch skip để tránh vòng lặp vô tận
+    total_batches = len(val_loader)
+    processed_batches = 0
+    
+    with torch.no_grad():
+        for batch_idx, (src, tgt, src_len, tgt_len) in enumerate(tqdm(val_loader, desc='Validation', ncols=120, file=None, dynamic_ncols=False)):
+            # Kiểm tra nếu đã skip quá nhiều - dừng validation để tránh vòng lặp vô tận
+            if skipped_batches >= max_skipped_batches:
+                print(f"\n⚠️  Đã skip {skipped_batches} batches trong validation (giới hạn: {max_skipped_batches})")
+                print("   Dừng validation để tránh vòng lặp vô tận")
+                break
+            
+            # Kiểm tra nếu đã xử lý đủ batches (ít nhất 50% để có kết quả hợp lý)
+            if processed_batches > 0 and skipped_batches > 0:
+                success_rate = processed_batches / (processed_batches + skipped_batches)
+                if success_rate < 0.3:  # Nếu tỷ lệ thành công < 30%
+                    print(f"\n⚠️  Tỷ lệ thành công quá thấp ({success_rate*100:.1f}%)")
+                    print("   Dừng validation để tránh kết quả không chính xác")
+                    break
+            
+            try:
+                # Kiểm tra memory spike
+                if device.type == 'cuda' and batch_idx % 50 == 0:
+                    is_spike, mem_info = check_memory_spike(threshold=0.9)
+                    if is_spike and mem_info:
+                        print(f"\n⚠️  GPU Memory spike trong validation: {mem_info['usage_percent']:.1f}%")
+                        clear_gpu_cache()
+                
+                # Kiểm tra batch size quá lớn
+                max_seq_len = max(src.size(1), tgt.size(1))
+                if max_seq_len > 200:
+                    skipped_batches += 1
+                    if skipped_batches % 5 == 0:
+                        print(f"\n⚠️  Đã skip {skipped_batches} batches do sequence quá dài (>200 tokens)")
+                    continue
+                
+                # Move to device
+                src = src.to(device)
+                tgt = tgt.to(device)
+                
+                # Target input và output
+                tgt_input = tgt[:, :-1]
+                tgt_output = tgt[:, 1:]
+                
+                # Forward pass với dtype phù hợp (API mới torch.amp.autocast)
+                if use_bf16:
+                    with torch.amp.autocast('cuda', enabled=True, dtype=torch.bfloat16):
+                        output = model(src, tgt_input)
+                        loss = criterion(output, tgt_output)
+                elif use_amp:
+                    with torch.amp.autocast('cuda', enabled=True):
+                        output = model(src, tgt_input)
+                        loss = criterion(output, tgt_output)
+                else:
+                    output = model(src, tgt_input)
+                    loss = criterion(output, tgt_output)
+                
+                # Statistics
+                num_tokens = (tgt_output != criterion.pad_idx).sum().item()
+                total_loss += loss.item() * num_tokens
+                total_tokens += num_tokens
+                processed_batches += 1
+                
+                # Update progress log every 100 batches
+                current_loss = loss.item()
+                current_ppl = calculate_perplexity(current_loss)
+                if (batch_idx + 1) % 100 == 0:
+                    tqdm.write(f"Validation Batch {batch_idx+1}/{total_batches}: loss={current_loss:.4f}, ppl={current_ppl:.2f}")
+                
+            except RuntimeError as e:
+                # Xử lý OOM trong validation
+                error_msg = str(e)
+                if 'out of memory' in error_msg.lower() or 'cuda' in error_msg.lower():
+                    print(f"\n❌ OOM Error trong validation tại batch {batch_idx}")
+                    clear_gpu_cache()
+                    skipped_batches += 1
+                    if skipped_batches >= max_skipped_batches:
+                        print(f"\n⚠️  Đã skip {skipped_batches} batches trong validation (giới hạn: {max_skipped_batches})")
+                        print("   Dừng validation để tránh vòng lặp vô tận")
+                        break
+                    continue
+                else:
+                    raise
+    
+    # Kiểm tra kết quả
+    if total_tokens > 0:
+        avg_loss = total_loss / total_tokens
+        avg_perplexity = calculate_perplexity(avg_loss)
+    else:
+        avg_loss = float('inf')
+        avg_perplexity = float('inf')
+        print("\n⚠️  Không có batch nào được xử lý thành công trong validation!")
+    
+    if skipped_batches > 0:
+        print(f"\n⚠️  Đã skip {skipped_batches}/{total_batches} batches trong validation")
+        print(f"   Đã xử lý thành công: {processed_batches} batches")
+    
+    return avg_loss, avg_perplexity
+
+# ============================================================================
+# 6. TRAINING LOOP
+# ============================================================================
+
+def train_model(
+    model,
+    train_loader,
+    val_loader,
+    num_epochs,
+    device,
+    d_model,
+    warmup_steps=4000,
+    label_smoothing=0.1,
+    grad_accum_steps=1,
+    precision="amp",
+    checkpoint_dir="checkpoints",
+    save_every=1,
+    resume_from=None,
+    target_max_lr: float | None = None,
+):
+    """
+    Huấn luyện model hoàn chỉnh
+    Args:
+        model: Transformer model
+        train_loader: Training DataLoader
+        val_loader: Validation DataLoader
+        num_epochs: Số epochs
+        device: Device (cuda/cpu)
+        d_model: Model dimension (để setup scheduler)
+        warmup_steps: Warmup steps cho scheduler (mặc định: 4000)
+        label_smoothing: Label smoothing factor
+        checkpoint_dir: Directory để lưu checkpoints
+        save_every: Lưu checkpoint mỗi N epochs
+        resume_from: Path to checkpoint to resume from (optional)
+        
+    Returns:
+        history: Dictionary chứa training history
+    """
+    use_amp = precision == 'amp' and device.type == 'cuda'
+    use_bf16 = precision == 'bf16' and device.type == 'cuda'
+    # Chỉ dùng GradScaler cho float16, bf16 không cần
+    scaler = GradScaler(enabled=use_amp and not use_bf16) if use_amp and not use_bf16 else None
+
+    checkpoint_dir = Path(checkpoint_dir)
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Setup optimizer
+    optimizer = optim.Adam(
+        model.parameters(),
+        lr=1.0,  # Sẽ được điều chỉnh bởi scheduler
+        betas=(0.9, 0.98),
+        eps=1e-9,
+    )
+    
+    # Setup scheduler
+    # Nếu người dùng chỉ định target_max_lr, tính factor tương ứng để đạt
+    # đỉnh LR xấp xỉ target_max_lr tại cuối warmup.
+    if target_max_lr is not None:
+        # lr_peak ≈ factor * d_model^(-0.5) * warmup_steps^(-0.5)
+        base = (d_model ** -0.5) * (warmup_steps ** -0.5)
+        factor = float(target_max_lr) / float(base)
+    else:
+        # Giữ nguyên behaviour cũ: peak LR ≈ 5e-4 cho d_model=384, warmup=10000
+        factor = 0.98
+
+    scheduler = TransformerLRScheduler(
+        optimizer,
+        d_model=d_model,
+        warmup_steps=warmup_steps,
+        factor=factor,
+    )
+    
+    # Setup loss function
+    criterion = LabelSmoothingLoss(
+        vocab_size=model.output_layer.out_features if hasattr(model, 'output_layer') else model.decoder.fc_out.out_features,
+        pad_idx=model.pad_idx,
+        smoothing=label_smoothing
+    )
+    
+    # Training history
+    history = {
+        'train_loss': [],
+        'train_ppl': [],
+        'val_loss': [],
+        'val_ppl': [],
+        'lr': []
+    }
+    
+    best_val_loss = float('inf')
+    start_epoch = 1
+    
+    # Resume from checkpoint if provided
+    if resume_from and Path(resume_from).exists():
+        print(f"\n→ Đang tiếp tục từ checkpoint: {resume_from}")
+        checkpoint = torch.load(resume_from, map_location=device)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+        start_epoch = checkpoint['epoch'] + 1
+        history = checkpoint.get('history', history)
+        best_val_loss = min(history.get('val_loss', [float('inf')]))
+        print(f"→ Tiếp tục từ epoch {start_epoch}")
+        print(f"→ Best validation loss trước đó: {best_val_loss:.4f}")
+    
+    print("="*70)
+    print("BẮT ĐẦU HUẤN LUYỆN")
+    print("="*70)
+    print(f"Device: {device}")
+    print(f"Number of epochs: {num_epochs}")
+    print(f"Starting from epoch: {start_epoch}")
+    print(f"Steps per epoch: {len(train_loader):,}")
+    print(f"Warmup steps: {warmup_steps:,}")
+    print(f"Label smoothing: {label_smoothing}")
+    print(f"Grad accumulation steps: {grad_accum_steps}")
+    if use_bf16:
+        print(f"Precision: BF16 (bfloat16 - recommended for stability)")
+    elif use_amp:
+        print(f"Precision: AMP (float16 mixed precision)")
+    else:
+        print(f"Precision: FP32")
+    print("="*70 + "\n")
+    
+    for epoch in range(start_epoch, num_epochs + 1):
+        epoch_start_time = time.time()
+        
+        # Training
+        train_loss, train_ppl = train_epoch(
+            model,
+            train_loader,
+            optimizer,
+            scheduler,
+            criterion,
+            device,
+            epoch,
+            grad_accum_steps=grad_accum_steps,
+            use_amp=use_amp,
+            use_bf16=use_bf16,
+            scaler=scaler,
+        )
+        
+        # Validation
+        val_loss, val_ppl = validate(
+            model,
+            val_loader,
+            criterion,
+            device,
+            use_amp=use_amp,
+            use_bf16=use_bf16,
+        )
+        
+        epoch_time = time.time() - epoch_start_time
+        
+        # Lưu history
+        history['train_loss'].append(train_loss)
+        history['train_ppl'].append(train_ppl)
+        history['val_loss'].append(val_loss)
+        history['val_ppl'].append(val_ppl)
+        history['lr'].append(scheduler.get_last_lr()[0])
+        
+        # In kết quả - single line format
+        checkpoint_msg = ""
+        best_msg = ""
+        
+        # Lưu checkpoint
+        if epoch % save_every == 0:
+            checkpoint_path = checkpoint_dir / f'checkpoint_epoch_{epoch}.pt'
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'scheduler_state_dict': scheduler.state_dict(),
+                'train_loss': train_loss,
+                'val_loss': val_loss,
+                'history': history
+            }, checkpoint_path)
+            checkpoint_msg = " | ✓ Saved checkpoint"
+        
+        # Lưu best model
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            best_model_path = checkpoint_dir / 'best_model.pt'
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'val_loss': val_loss,
+                'val_ppl': val_ppl
+            }, best_model_path)
+            best_msg = " | ✓ New best model!"
+        
+        print(f"Epoch {epoch}/{num_epochs} | Time: {epoch_time:.2f}s | Train Loss: {train_loss:.4f} | Train PPL: {train_ppl:.2f} | Val Loss: {val_loss:.4f} | Val PPL: {val_ppl:.2f} | LR: {scheduler.get_last_lr()[0]:.6f}{checkpoint_msg}{best_msg}")
+    
+    print("\n" + "="*70)
+    print("HOÀN TẤT HUẤN LUYỆN")
+    print("="*70)
+    print(f"Best validation loss: {best_val_loss:.4f}")
+    print(f"Best validation perplexity: {calculate_perplexity(best_val_loss):.2f}")
+    
+    return history
+
+# ============================================================================
+# 7. PLOT TRAINING HISTORY
+# ============================================================================
+
+def plot_training_history(history, save_path='training_history.png'):
+    """
+    Vẽ đồ thị training history
+    
+    Args:
+        history: Training history dictionary
+        save_path: Path để lưu hình
+    """
+    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
+    
+    epochs = range(1, len(history['train_loss']) + 1)
+    
+    # Plot Loss
+    axes[0].plot(epochs, history['train_loss'], 'b-', label='Train Loss')
+    axes[0].plot(epochs, history['val_loss'], 'r-', label='Val Loss')
+    axes[0].set_xlabel('Epoch')
+    axes[0].set_ylabel('Loss')
+    axes[0].set_title('Training and Validation Loss')
+    axes[0].legend()
+    axes[0].grid(True)
+    
+    # Plot Perplexity
+    axes[1].plot(epochs, history['train_ppl'], 'b-', label='Train PPL')
+    axes[1].plot(epochs, history['val_ppl'], 'r-', label='Val PPL')
+    axes[1].set_xlabel('Epoch')
+    axes[1].set_ylabel('Perplexity')
+    axes[1].set_title('Training and Validation Perplexity')
+    axes[1].legend()
+    axes[1].grid(True)
+    
+    # Plot Learning Rate
+    axes[2].plot(epochs, history['lr'], 'g-')
+    axes[2].set_xlabel('Epoch')
+    axes[2].set_ylabel('Learning Rate')
+    axes[2].set_title('Learning Rate Schedule')
+    axes[2].grid(True)
+    
+    save_path = Path(save_path)
+    save_path.parent.mkdir(parents=True, exist_ok=True)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"\n✓ Saved training history plot to {save_path}")
+    plt.close()
+
+# ============================================================================
+# 8. LOAD CHECKPOINT
+# ============================================================================
+
+def load_checkpoint(model, checkpoint_path, device):
+    """
+    Load checkpoint
+    
+    Args:
+        model: Transformer model
+        checkpoint_path: Path to checkpoint
+        device: Device
+        
+    Returns:
+        model: Model with loaded weights
+        checkpoint: Checkpoint dictionary
+    """
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    print(f"✓ Loaded checkpoint from {checkpoint_path}")
+    print(f"  Epoch: {checkpoint.get('epoch', 'N/A')}")
+    print(f"  Val Loss: {checkpoint.get('val_loss', 'N/A'):.4f}")
+    
+    return model, checkpoint

src/transformer_components.py

@@ -0,0 +1,378 @@
+"""
+PHẦN B: XÂY DỰNG KIẾN TRÚC TRANSFORMER FROM SCRATCH
+Các thành phần cốt lõi: Attention, Positional Encoding, Encoder, Decoder
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+# ============================================================================
+# 1. SCALED DOT-PRODUCT ATTENTION
+# ============================================================================
+
+class ScaledDotProductAttention(nn.Module):
+    """
+    Scaled Dot-Product Attention
+    
+    Attention(Q, K, V) = softmax(Q·K^T / sqrt(d_k)) · V
+    
+    Args:
+        d_k: Dimension của key (để scale)
+    """
+    def __init__(self, d_k):
+        super().__init__()
+        self.d_k = d_k
+        self.softmax = nn.Softmax(dim=-1)
+        
+    def forward(self, Q, K, V, mask=None):
+        """
+        Args:
+            Q: Query [batch_size, n_heads, seq_len, d_k]
+            K: Key [batch_size, n_heads, seq_len, d_k]
+            V: Value [batch_size, n_heads, seq_len, d_v]
+            mask: Mask [batch_size, 1, seq_len, seq_len] hoặc [batch_size, 1, 1, seq_len]
+            
+        Returns:
+            output: [batch_size, n_heads, seq_len, d_v]
+            attention_weights: [batch_size, n_heads, seq_len, seq_len]
+        """
+        # Tính attention scores: Q·K^T
+        # [batch, n_heads, seq_len, d_k] @ [batch, n_heads, d_k, seq_len]
+        # -> [batch, n_heads, seq_len, seq_len]
+        scores = torch.matmul(Q, K.transpose(-2, -1))
+        
+        # Scale bởi sqrt(d_k) để tránh gradient quá nhỏ
+        scores = scores / math.sqrt(self.d_k)
+        
+        # Apply mask (nếu có)
+        if mask is not None:
+            mask = mask.to(dtype=torch.bool)
+            neg_inf = torch.finfo(scores.dtype).min
+            scores = scores.masked_fill(~mask, neg_inf)
+        
+        # Softmax để có attention weights
+        attention_weights = self.softmax(scores)
+        
+        # Apply attention weights to values
+        # [batch, n_heads, seq_len, seq_len] @ [batch, n_heads, seq_len, d_v]
+        # -> [batch, n_heads, seq_len, d_v]
+        output = torch.matmul(attention_weights, V)
+        
+        return output, attention_weights
+
+# ============================================================================
+# 2. MULTI-HEAD ATTENTION
+# ============================================================================
+
+class MultiHeadAttention(nn.Module):
+    """
+    Multi-Head Attention
+    
+    Chia input thành nhiều heads, mỗi head học các representation khác nhau
+    
+    Args:
+        d_model: Dimension của model
+        n_heads: Số lượng attention heads
+        dropout: Dropout rate
+    """
+    def __init__(self, d_model, n_heads, dropout=0.1):
+        super().__init__()
+        
+        assert d_model % n_heads == 0, "d_model phải chia hết cho n_heads"
+        
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.d_k = d_model // n_heads  # Dimension của mỗi head
+        
+        # Linear layers để project Q, K, V
+        self.W_q = nn.Linear(d_model, d_model)
+        self.W_k = nn.Linear(d_model, d_model)
+        self.W_v = nn.Linear(d_model, d_model)
+        
+        # Scaled Dot-Product Attention
+        self.attention = ScaledDotProductAttention(self.d_k)
+        
+        # Output projection
+        self.W_o = nn.Linear(d_model, d_model)
+        
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, Q, K, V, mask=None):
+        """
+        Args:
+            Q: Query [batch_size, seq_len, d_model]
+            K: Key [batch_size, seq_len, d_model]
+            V: Value [batch_size, seq_len, d_model]
+            mask: Mask tensor
+            
+        Returns:
+            output: [batch_size, seq_len, d_model]
+            attention_weights: [batch_size, n_heads, seq_len, seq_len]
+        """
+        batch_size = Q.size(0)
+        
+        # 1. Linear projection và chia thành multiple heads
+        # [batch, seq_len, d_model] -> [batch, seq_len, n_heads, d_k] -> [batch, n_heads, seq_len, d_k]
+        Q = self.W_q(Q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
+        K = self.W_k(K).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
+        V = self.W_v(V).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
+        
+        # 2. Apply attention
+        # output: [batch, n_heads, seq_len, d_k]
+        output, attention_weights = self.attention(Q, K, V, mask)
+        
+        # 3. Concatenate heads
+        # [batch, n_heads, seq_len, d_k] -> [batch, seq_len, n_heads, d_k] -> [batch, seq_len, d_model]
+        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
+        
+        # 4. Final linear projection
+        output = self.W_o(output)
+        output = self.dropout(output)
+        
+        return output, attention_weights
+
+# ============================================================================
+# 3. POSITION-WISE FEED-FORWARD NETWORK
+# ============================================================================
+
+class PositionwiseFeedForward(nn.Module):
+    """
+    Position-wise Feed-Forward Network
+    
+    FFN(x) = max(0, xW1 + b1)W2 + b2
+    
+    Áp dụng 2 linear transformations với ReLU ở giữa
+    
+    Args:
+        d_model: Dimension của model
+        d_ff: Dimension của hidden layer (thường = 4 * d_model)
+        dropout: Dropout rate
+    """
+    def __init__(self, d_model, d_ff, dropout=0.1):
+        super().__init__()
+        
+        self.linear1 = nn.Linear(d_model, d_ff)
+        self.linear2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, x):
+        """
+        Args:
+            x: [batch_size, seq_len, d_model]
+            
+        Returns:
+            output: [batch_size, seq_len, d_model]
+        """
+        # x -> W1 -> ReLU -> Dropout -> W2 -> Dropout
+        x = self.linear1(x)
+        x = F.relu(x)
+        x = self.dropout(x)
+        x = self.linear2(x)
+        x = self.dropout(x)
+        
+        return x
+
+# ============================================================================
+# 4. POSITIONAL ENCODING
+# ============================================================================
+
+class PositionalEncoding(nn.Module):
+    """
+    Positional Encoding (Sinusoidal)
+    
+    PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
+    PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
+    
+    Thêm thông tin về vị trí của token trong sequence
+    
+    Args:
+        d_model: Dimension của model
+        max_len: Maximum sequence length
+        dropout: Dropout rate
+    """
+    def __init__(self, d_model, max_len=5000, dropout=0.1):
+        super().__init__()
+        
+        self.dropout = nn.Dropout(dropout)
+        
+        # Tạo positional encoding matrix
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len).unsqueeze(1).float()
+        
+        # Tính div_term cho công thức
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * 
+                            -(math.log(10000.0) / d_model))
+        
+        # Apply sin cho các vị trí chẵn, cos cho các vị trí lẻ
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        
+        # Thêm batch dimension
+        pe = pe.unsqueeze(0)  # [1, max_len, d_model]
+        
+        # Register as buffer (không train)
+        self.register_buffer('pe', pe)
+        
+    def forward(self, x):
+        """
+        Args:
+            x: [batch_size, seq_len, d_model]
+            
+        Returns:
+            output: [batch_size, seq_len, d_model]
+        """
+        seq_len = x.size(1)
+        
+        # Cộng positional encoding
+        x = x + self.pe[:, :seq_len, :]
+        
+        return self.dropout(x)
+
+# ============================================================================
+# 5. LAYER NORMALIZATION
+# ============================================================================
+
+class LayerNorm(nn.Module):
+    """
+    Layer Normalization
+    
+    Chuẩn hóa theo dimension cuối (features)
+    
+    Args:
+        d_model: Dimension của model
+        eps: Epsilon cho numerical stability
+    """
+    def __init__(self, d_model, eps=1e-6):
+        super().__init__()
+        
+        # Learnable parameters
+        self.gamma = nn.Parameter(torch.ones(d_model))
+        self.beta = nn.Parameter(torch.zeros(d_model))
+        self.eps = eps
+        
+    def forward(self, x):
+        """
+        Args:
+            x: [batch_size, seq_len, d_model]
+            
+        Returns:
+            output: [batch_size, seq_len, d_model]
+        """
+        mean = x.mean(-1, keepdim=True)
+        std = x.std(-1, keepdim=True)
+        
+        return self.gamma * (x - mean) / (std + self.eps) + self.beta
+
+# ============================================================================
+# 6. RESIDUAL CONNECTION
+# ============================================================================
+
+class ResidualConnection(nn.Module):
+    """
+    Residual Connection với Layer Normalization
+    
+    output = LayerNorm(x + Sublayer(x))
+    
+    Args:
+        d_model: Dimension của model
+        dropout: Dropout rate
+    """
+    def __init__(self, d_model, dropout=0.1):
+        super().__init__()
+        
+        self.norm = LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, x, sublayer):
+        """
+        Args:
+            x: Input [batch_size, seq_len, d_model]
+            sublayer: Function (callable) để apply
+            
+        Returns:
+            output: [batch_size, seq_len, d_model]
+        """
+        # Residual: x + sublayer(norm(x))
+        return x + self.dropout(sublayer(self.norm(x)))
+
+# ============================================================================
+# 7. EMBEDDING LAYER
+# ============================================================================
+
+class Embedding(nn.Module):
+    """
+    Embedding layer với scaling
+    
+    Args:
+        vocab_size: Kích thước vocabulary
+        d_model: Dimension của model
+    """
+    def __init__(self, vocab_size, d_model):
+        super().__init__()
+        
+        self.d_model = d_model
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        
+    def forward(self, x):
+        """
+        Args:
+            x: [batch_size, seq_len]
+            
+        Returns:
+            output: [batch_size, seq_len, d_model]
+        """
+        # Scale embedding by sqrt(d_model) như trong paper
+        return self.embedding(x) * math.sqrt(self.d_model)
+
+# ============================================================================
+# 8. TEST COMPONENTS
+# ============================================================================
+
+if __name__ == "__main__":
+    print("="*70)
+    print("KIỂM TRA CÁC THÀNH PHẦN TRANSFORMER")
+    print("="*70)
+    
+    # Hyperparameters
+    batch_size = 2
+    seq_len = 10
+    d_model = 512
+    n_heads = 8
+    d_ff = 2048
+    vocab_size = 10000
+    
+    # Test Embedding
+    print("\n1. Test Embedding:")
+    embedding = Embedding(vocab_size, d_model)
+    x = torch.randint(0, vocab_size, (batch_size, seq_len))
+    embedded = embedding(x)
+    print(f"   Input shape: {x.shape}")
+    print(f"   Output shape: {embedded.shape}")
+    
+    # Test Positional Encoding
+    print("\n2. Test Positional Encoding:")
+    pos_enc = PositionalEncoding(d_model)
+    pos_encoded = pos_enc(embedded)
+    print(f"   Input shape: {embedded.shape}")
+    print(f"   Output shape: {pos_encoded.shape}")
+    
+    # Test Multi-Head Attention
+    print("\n3. Test Multi-Head Attention:")
+    mha = MultiHeadAttention(d_model, n_heads)
+    output, attn_weights = mha(pos_encoded, pos_encoded, pos_encoded)
+    print(f"   Input shape: {pos_encoded.shape}")
+    print(f"   Output shape: {output.shape}")
+    print(f"   Attention weights shape: {attn_weights.shape}")
+    
+    # Test Feed-Forward Network
+    print("\n4. Test Feed-Forward Network:")
+    ffn = PositionwiseFeedForward(d_model, d_ff)
+    ffn_output = ffn(output)
+    print(f"   Input shape: {output.shape}")
+    print(f"   Output shape: {ffn_output.shape}")
+    
+    print("\n" + "="*70)
+    print("✓ TẤT CẢ THÀNH PHẦN HOẠT ĐỘNG ĐÚNG!")
+    print("="*70)

src/transformer_encoder_decoder.py

@@ -0,0 +1,403 @@
+"""
+TRANSFORMER ENCODER & DECODER (FIXED)
+Xây dựng hoàn chỉnh Encoder và Decoder layers
+"""
+
+import torch
+import torch.nn as nn
+from .transformer_components import (
+    MultiHeadAttention,
+    PositionwiseFeedForward,
+    ResidualConnection,
+    LayerNorm
+)
+
+# ============================================================================
+# 1. ENCODER LAYER
+# ============================================================================
+
+class EncoderLayer(nn.Module):
+    """
+    Một layer của Transformer Encoder
+    
+    Gồm:
+    1. Multi-Head Self-Attention
+    2. Add & Norm
+    3. Feed-Forward Network
+    4. Add & Norm
+    
+    Args:
+        d_model: Dimension của model
+        n_heads: Số lượng attention heads
+        d_ff: Dimension của feed-forward network
+        dropout: Dropout rate
+    """
+    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
+        super().__init__()
+        
+        # Multi-Head Self-Attention
+        self.self_attention = MultiHeadAttention(d_model, n_heads, dropout)
+        
+        # Feed-Forward Network
+        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
+        
+        # Residual Connections
+        self.residual1 = ResidualConnection(d_model, dropout)
+        self.residual2 = ResidualConnection(d_model, dropout)
+        
+    def forward(self, x, mask=None):
+        """
+        Args:
+            x: Input [batch_size, seq_len, d_model]
+            mask: Mask tensor [batch_size, 1, 1, seq_len] (để mask padding)
+            
+        Returns:
+            output: [batch_size, seq_len, d_model]
+        """
+        # 1. Self-Attention với Residual Connection
+        x = self.residual1(x, lambda x: self.self_attention(x, x, x, mask)[0])
+        
+        # 2. Feed-Forward với Residual Connection
+        x = self.residual2(x, self.feed_forward)
+        
+        return x
+
+# ============================================================================
+# 2. ENCODER
+# ============================================================================
+
+class Encoder(nn.Module):
+    """
+    Transformer Encoder - Stack của N encoder layers
+    
+    Args:
+        vocab_size: Kích thước vocabulary
+        d_model: Dimension của model
+        n_layers: Số lượng encoder layers
+        n_heads: Số lượng attention heads
+        d_ff: Dimension của feed-forward network
+        dropout: Dropout rate
+        max_len: Maximum sequence length
+    """
+    def __init__(self, vocab_size, d_model, n_layers, n_heads, d_ff, dropout=0.1, max_len=5000):
+        super().__init__()
+        
+        from .transformer_components import Embedding, PositionalEncoding
+        
+        # Embedding layer
+        self.embedding = Embedding(vocab_size, d_model)
+        
+        # Positional Encoding
+        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)
+        
+        # Stack of Encoder Layers
+        self.layers = nn.ModuleList([
+            EncoderLayer(d_model, n_heads, d_ff, dropout)
+            for _ in range(n_layers)
+        ])
+        
+        # Final Layer Normalization
+        self.norm = LayerNorm(d_model)
+        
+    def forward(self, src, src_mask=None):
+        """
+        Args:
+            src: Source sequence [batch_size, src_len]
+            src_mask: Source mask [batch_size, 1, 1, src_len]
+            
+        Returns:
+            output: [batch_size, src_len, d_model]
+        """
+        # 1. Embedding + Positional Encoding
+        x = self.embedding(src)
+        x = self.pos_encoding(x)
+        
+        # 2. Pass through encoder layers
+        for layer in self.layers:
+            x = layer(x, src_mask)
+        
+        # 3. Final normalization
+        x = self.norm(x)
+        
+        return x
+
+# ============================================================================
+# 3. DECODER LAYER
+# ============================================================================
+
+class DecoderLayer(nn.Module):
+    """
+    Một layer của Transformer Decoder
+    
+    Gồm:
+    1. Masked Multi-Head Self-Attention
+    2. Add & Norm
+    3. Multi-Head Cross-Attention (với Encoder output)
+    4. Add & Norm
+    5. Feed-Forward Network
+    6. Add & Norm
+    
+    Args:
+        d_model: Dimension của model
+        n_heads: Số lượng attention heads
+        d_ff: Dimension của feed-forward network
+        dropout: Dropout rate
+    """
+    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
+        super().__init__()
+        
+        # Masked Multi-Head Self-Attention
+        self.self_attention = MultiHeadAttention(d_model, n_heads, dropout)
+        
+        # Multi-Head Cross-Attention (Encoder-Decoder Attention)
+        self.cross_attention = MultiHeadAttention(d_model, n_heads, dropout)
+        
+        # Feed-Forward Network
+        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
+        
+        # Residual Connections
+        self.residual1 = ResidualConnection(d_model, dropout)
+        self.residual2 = ResidualConnection(d_model, dropout)
+        self.residual3 = ResidualConnection(d_model, dropout)
+        
+    def forward(self, x, encoder_output, src_mask=None, tgt_mask=None):
+        """
+        Args:
+            x: Target input [batch_size, tgt_len, d_model]
+            encoder_output: Encoder output [batch_size, src_len, d_model]
+            src_mask: Source mask [batch_size, 1, 1, src_len]
+            tgt_mask: Target mask [batch_size, 1, tgt_len, tgt_len] (causal mask)
+            
+        Returns:
+            output: [batch_size, tgt_len, d_model]
+        """
+        # 1. Masked Self-Attention với Residual Connection
+        x = self.residual1(x, lambda x: self.self_attention(x, x, x, tgt_mask)[0])
+        
+        # 2. Cross-Attention với Encoder output
+        # Q từ decoder, K, V từ encoder
+        x = self.residual2(x, lambda x: self.cross_attention(x, encoder_output, encoder_output, src_mask)[0])
+        
+        # 3. Feed-Forward với Residual Connection
+        x = self.residual3(x, self.feed_forward)
+        
+        return x
+
+# ============================================================================
+# 4. DECODER
+# ============================================================================
+
+class Decoder(nn.Module):
+    """
+    Transformer Decoder - Stack của N decoder layers
+    
+    Args:
+        vocab_size: Kích thước vocabulary
+        d_model: Dimension của model
+        n_layers: Số lượng decoder layers
+        n_heads: Số lượng attention heads
+        d_ff: Dimension của feed-forward network
+        dropout: Dropout rate
+        max_len: Maximum sequence length
+    """
+    def __init__(self, vocab_size, d_model, n_layers, n_heads, d_ff, dropout=0.1, max_len=5000):
+        super().__init__()
+        
+        from .transformer_components import Embedding, PositionalEncoding
+        
+        # Embedding layer
+        self.embedding = Embedding(vocab_size, d_model)
+        
+        # Positional Encoding
+        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)
+        
+        # Stack of Decoder Layers
+        self.layers = nn.ModuleList([
+            DecoderLayer(d_model, n_heads, d_ff, dropout)
+            for _ in range(n_layers)
+        ])
+        
+        # Final Layer Normalization
+        self.norm = LayerNorm(d_model)
+        
+        # Output projection to vocabulary
+        self.fc_out = nn.Linear(d_model, vocab_size)
+        
+    def forward(self, tgt, encoder_output, src_mask=None, tgt_mask=None):
+        """
+        Args:
+            tgt: Target sequence [batch_size, tgt_len]
+            encoder_output: Encoder output [batch_size, src_len, d_model]
+            src_mask: Source mask [batch_size, 1, 1, src_len]
+            tgt_mask: Target mask [batch_size, 1, tgt_len, tgt_len]
+            
+        Returns:
+            output: [batch_size, tgt_len, vocab_size]
+        """
+        # 1. Embedding + Positional Encoding
+        x = self.embedding(tgt)
+        x = self.pos_encoding(x)
+        
+        # 2. Pass through decoder layers
+        for layer in self.layers:
+            x = layer(x, encoder_output, src_mask, tgt_mask)
+        
+        # 3. Final normalization
+        x = self.norm(x)
+        
+        # 4. Project to vocabulary
+        output = self.fc_out(x)
+        
+        return output
+
+# ============================================================================
+# 5. MASK FUNCTIONS (FIXED)
+# ============================================================================
+
+def create_padding_mask(seq, pad_idx=0):
+    """
+    Tạo mask cho padding tokens
+    
+    Args:
+        seq: Sequence [batch_size, seq_len]
+        pad_idx: Index của padding token
+        
+    Returns:
+        mask: [batch_size, 1, 1, seq_len] (bool type)
+    """
+    # Tạo mask: True cho non-padding, False cho padding
+    mask = (seq != pad_idx).unsqueeze(1).unsqueeze(2)
+    return mask  # Returns bool tensor
+
+def create_causal_mask(seq_len, device):
+    """
+    Tạo causal mask (look-ahead mask) cho decoder
+    Ngăn decoder nhìn thấy future tokens
+    
+    Args:
+        seq_len: Length của sequence
+        device: Device (cuda hoặc cpu)
+        
+    Returns:
+        mask: [1, 1, seq_len, seq_len] (bool type)
+    """
+    # Tạo lower triangular matrix - FIXED: convert to bool
+    mask = torch.tril(torch.ones(seq_len, seq_len, device=device))
+    mask = mask.bool()  # Convert to bool
+    mask = mask.unsqueeze(0).unsqueeze(1)
+    return mask
+
+def create_target_mask(tgt, pad_idx=0):
+    """
+    Tạo mask kết hợp cho target sequence (padding + causal)
+    
+    Args:
+        tgt: Target sequence [batch_size, tgt_len]
+        pad_idx: Index của padding token
+        
+    Returns:
+        mask: [batch_size, 1, tgt_len, tgt_len] (bool type)
+    """
+    batch_size, tgt_len = tgt.size()
+    device = tgt.device
+    
+    # Padding mask - returns bool
+    padding_mask = (tgt != pad_idx).unsqueeze(1).unsqueeze(2)  # [batch, 1, 1, tgt_len]
+    
+    # Causal mask - returns bool
+    causal_mask = create_causal_mask(tgt_len, device)  # [1, 1, tgt_len, tgt_len]
+    
+    # Kết hợp cả 2 masks - both are bool now
+    mask = padding_mask & causal_mask
+    
+    return mask
+
+# ============================================================================
+# 6. TEST ENCODER & DECODER
+# ============================================================================
+
+if __name__ == "__main__":
+    print("="*70)
+    print("KIỂM TRA ENCODER & DECODER")
+    print("="*70)
+    
+    # Hyperparameters
+    batch_size = 2
+    src_len = 10
+    tgt_len = 12
+    src_vocab_size = 10000
+    tgt_vocab_size = 8000
+    d_model = 512
+    n_layers = 6
+    n_heads = 8
+    d_ff = 2048
+    dropout = 0.1
+    pad_idx = 0
+    
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"\nDevice: {device}")
+    
+    # Tạo dummy data
+    src = torch.randint(1, src_vocab_size, (batch_size, src_len)).to(device)
+    tgt = torch.randint(1, tgt_vocab_size, (batch_size, tgt_len)).to(device)
+    
+    # Tạo masks
+    src_mask = create_padding_mask(src, pad_idx).to(device)
+    tgt_mask = create_target_mask(tgt, pad_idx).to(device)
+    
+    print(f"\nInput shapes:")
+    print(f"  Source: {src.shape}")
+    print(f"  Target: {tgt.shape}")
+    print(f"  Source mask: {src_mask.shape}, dtype: {src_mask.dtype}")
+    print(f"  Target mask: {tgt_mask.shape}, dtype: {tgt_mask.dtype}")
+    
+    # Test Encoder
+    print("\n" + "="*70)
+    print("Test Encoder")
+    print("="*70)
+    
+    encoder = Encoder(
+        vocab_size=src_vocab_size,
+        d_model=d_model,
+        n_layers=n_layers,
+        n_heads=n_heads,
+        d_ff=d_ff,
+        dropout=dropout
+    ).to(device)
+    
+    encoder_output = encoder(src, src_mask)
+    print(f"Encoder output shape: {encoder_output.shape}")
+    print(f"Expected: [{batch_size}, {src_len}, {d_model}]")
+    
+    # Test Decoder
+    print("\n" + "="*70)
+    print("Test Decoder")
+    print("="*70)
+    
+    decoder = Decoder(
+        vocab_size=tgt_vocab_size,
+        d_model=d_model,
+        n_layers=n_layers,
+        n_heads=n_heads,
+        d_ff=d_ff,
+        dropout=dropout
+    ).to(device)
+    
+    decoder_output = decoder(tgt, encoder_output, src_mask, tgt_mask)
+    print(f"Decoder output shape: {decoder_output.shape}")
+    print(f"Expected: [{batch_size}, {tgt_len}, {tgt_vocab_size}]")
+    
+    # Số lượng parameters
+    encoder_params = sum(p.numel() for p in encoder.parameters())
+    decoder_params = sum(p.numel() for p in decoder.parameters())
+    
+    print("\n" + "="*70)
+    print("THỐNG KÊ MÔ HÌNH")
+    print("="*70)
+    print(f"Encoder parameters: {encoder_params:,}")
+    print(f"Decoder parameters: {decoder_params:,}")
+    print(f"Total parameters: {encoder_params + decoder_params:,}")
+    
+    print("\n" + "="*70)
+    print("✓ ENCODER & DECODER HOẠT ĐỘNG ĐÚNG!")
+    print("="*70)