src/app.py

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