inference.py

"""
Inference Pipeline for Code-Mixed Machine Translation
Supports: batch translation, ONNX export, and quantization
"""

import os
import time
import logging
from typing import List, Optional, Union, Dict

import torch
import numpy as np
from transformers import (
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    pipeline,
    GenerationConfig,
)

logger = logging.getLogger(__name__)


# ─────────────────────────────────────────────────────────────────────────────
# 1. CORE INFERENCE PIPELINE
# ─────────────────────────────────────────────────────────────────────────────

class CodeMixedTranslator:
    """
    Production inference pipeline for code-mixed MT using NLLB-200.
    Supports single/batch translation with configurable generation parameters.
    """

    def __init__(
        self,
        model_path: str,
        src_lang: str = "hin_Deva",
        tgt_lang: str = "eng_Latn",
        device: Optional[str] = None,
        use_quantization: bool = False,
        max_length: int = 256,
    ):
        self.src_lang = src_lang
        self.tgt_lang = tgt_lang
        self.max_length = max_length
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        logger.info(f"Initializing translator: {src_lang} → {tgt_lang}")
        logger.info(f"Device: {self.device}")

        # Load tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_path,
            src_lang=src_lang,
            tgt_lang=tgt_lang,
        )

        # Load model with optional quantization
        if use_quantization and self.device == "cuda":
            logger.info("Loading with 8-bit quantization (bitsandbytes)...")
            try:
                from transformers import BitsAndBytesConfig
                quantization_config = BitsAndBytesConfig(load_in_8bit=True)
                self.model = AutoModelForSeq2SeqLM.from_pretrained(
                    model_path,
                    quantization_config=quantization_config,
                    device_map="auto",
                )
            except ImportError:
                logger.warning("bitsandbytes not installed. Loading without quantization.")
                self.model = self._load_standard_model(model_path)
        else:
            self.model = self._load_standard_model(model_path)

        self.model.eval()

        # Configure generation
        self.forced_bos_token_id = self.tokenizer.lang_code_to_id.get(tgt_lang)
        if self.forced_bos_token_id is None:
            raise ValueError(f"Language code '{tgt_lang}' not found in NLLB tokenizer.")

        logger.info(f"Translator ready. Forced BOS token ID: {self.forced_bos_token_id}")

    def _load_standard_model(self, model_path: str) -> AutoModelForSeq2SeqLM:
        """Standard model loading without quantization."""
        dtype = torch.float16 if self.device == "cuda" else torch.float32
        model = AutoModelForSeq2SeqLM.from_pretrained(
            model_path,
            torch_dtype=dtype,
        )
        model = model.to(self.device)
        return model

    def translate(
        self,
        texts: Union[str, List[str]],
        num_beams: int = 4,
        length_penalty: float = 1.0,
        no_repeat_ngram_size: int = 3,
        temperature: float = 1.0,
        do_sample: bool = False,
        return_scores: bool = False,
    ) -> Union[str, List[str], Dict]:
        """
        Translate code-mixed text(s) to target language.

        Args:
            texts: Single string or list of strings to translate
            num_beams: Beam search width (higher = better quality, slower)
            length_penalty: >1 encourages longer, <1 shorter translations
            no_repeat_ngram_size: Prevent n-gram repetition
            temperature: Sampling temperature (when do_sample=True)
            do_sample: Use sampling instead of beam search
            return_scores: Return generation scores alongside translations

        Returns:
            Translated string, list of strings, or dict with scores
        """
        single_input = isinstance(texts, str)
        if single_input:
            texts = [texts]

        # Tokenize
        self.tokenizer.src_lang = self.src_lang
        inputs = self.tokenizer(
            texts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=self.max_length,
        ).to(self.device)

        # Generate translations
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                forced_bos_token_id=self.forced_bos_token_id,
                num_beams=num_beams,
                max_length=self.max_length,
                length_penalty=length_penalty,
                no_repeat_ngram_size=no_repeat_ngram_size,
                temperature=temperature,
                do_sample=do_sample,
                return_dict_in_generate=return_scores,
                output_scores=return_scores,
                early_stopping=True,
            )

        if return_scores:
            sequences = outputs.sequences
            scores = outputs.sequences_scores.cpu().numpy().tolist() if hasattr(outputs, 'sequences_scores') else []
        else:
            sequences = outputs

        # Decode
        translations = self.tokenizer.batch_decode(
            sequences, skip_special_tokens=True
        )

        result = translations[0] if single_input else translations

        if return_scores:
            return {"translations": result, "scores": scores}

        return result

    def translate_batch_chunked(
        self,
        texts: List[str],
        batch_size: int = 16,
        **translation_kwargs,
    ) -> List[str]:
        """
        Translate a large list of texts in batches with progress tracking.
        Memory-efficient for large datasets.
        """
        all_translations = []
        total = len(texts)

        for i in range(0, total, batch_size):
            batch = texts[i: i + batch_size]
            start_time = time.time()

            translations = self.translate(batch, **translation_kwargs)
            if isinstance(translations, str):
                translations = [translations]

            all_translations.extend(translations)

            elapsed = time.time() - start_time
            logger.info(
                f"Batch {i // batch_size + 1}/{(total + batch_size - 1) // batch_size} | "
                f"{len(batch)} samples | {elapsed:.2f}s | "
                f"Speed: {len(batch) / elapsed:.1f} samples/s"
            )

        return all_translations

    def interactive_demo(self):
        """Run an interactive translation session in terminal."""
        print("\n" + "=" * 60)
        print(f"Code-Mixed MT: {self.src_lang} → {self.tgt_lang}")
        print("Type 'quit' to exit | 'lang' to show language codes")
        print("=" * 60 + "\n")

        while True:
            try:
                user_input = input("Enter text: ").strip()

                if user_input.lower() == 'quit':
                    print("Goodbye!")
                    break
                elif user_input.lower() == 'lang':
                    print(f"  Source: {self.src_lang}")
                    print(f"  Target: {self.tgt_lang}")
                    continue
                elif not user_input:
                    continue

                start = time.time()
                translation = self.translate(user_input)
                elapsed = time.time() - start

                print(f"Translation: {translation}")
                print(f"[{elapsed*1000:.0f}ms]\n")

            except KeyboardInterrupt:
                print("\nInterrupted. Goodbye!")
                break


# ─────────────────────────────────────────────────────────────────────────────
# 2. HUGGINGFACE PIPELINE WRAPPER
# ─────────────────────────────────────────────────────────────────────────────

def build_hf_pipeline(
    model_path: str,
    src_lang: str = "hin_Deva",
    tgt_lang: str = "eng_Latn",
    device: int = -1,   # -1 = CPU, 0+ = GPU index
    max_length: int = 256,
):
    """
    Build a standard HuggingFace translation pipeline.
    Easier to use for simple translation tasks.

    Args:
        model_path: Path to fine-tuned model or HuggingFace model name
        src_lang: NLLB source language code
        tgt_lang: NLLB target language code
        device: -1 for CPU, 0 for first GPU, etc.
        max_length: Maximum generation length

    Returns:
        HuggingFace pipeline for translation
    """
    tokenizer = AutoTokenizer.from_pretrained(model_path, src_lang=src_lang)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)

    forced_bos_token_id = tokenizer.lang_code_to_id[tgt_lang]

    translator = pipeline(
        task="translation",
        model=model,
        tokenizer=tokenizer,
        src_lang=src_lang,
        tgt_lang=tgt_lang,
        forced_bos_token_id=forced_bos_token_id,
        max_length=max_length,
        device=device,
    )

    return translator


# ─────────────────────────────────────────────────────────────────────────────
# 3. ONNX EXPORT (OPTIONAL - FOR DEPLOYMENT OPTIMIZATION)
# ─────────────────────────────────────────────────────────────────────────────

def export_to_onnx(model_path: str, output_dir: str, opset: int = 13):
    """
    Export fine-tuned NLLB model to ONNX format for optimized inference.
    Requires: optimum library (pip install optimum[onnxruntime])

    Args:
        model_path: Path to the fine-tuned HuggingFace model
        output_dir: Directory to save ONNX model
        opset: ONNX opset version
    """
    try:
        from optimum.onnxruntime import ORTModelForSeq2SeqLM
        from optimum.exporters.onnx import main_export

        os.makedirs(output_dir, exist_ok=True)
        logger.info(f"Exporting model to ONNX: {output_dir}")

        main_export(
            model_name_or_path=model_path,
            output=output_dir,
            task="text2text-generation",
            opset=opset,
        )

        logger.info(f"ONNX export complete. Files saved to: {output_dir}")
        return True

    except ImportError:
        logger.error("ONNX export requires: pip install optimum[onnxruntime]")
        return False


# ─────────────────────────────────────────────────────────────────────────────
# 4. TRANSLATION EXAMPLES
# ─────────────────────────────────────────────────────────────────────────────

def run_translation_examples(translator: CodeMixedTranslator, language_pair: str):
    """Run example translations to verify the model."""

    examples = {
        "hi-en": [
            "Mujhe bahut zyada hunger lag raha hai",
            "Aaj ka weather bahut acha hai, let's go for a walk",
            "Main office se late aaya, traffic bahut zyada tha",
            "Yeh movie bahut boring thi, waste of time",
            "Kya aap mujhe help kar sakte ho is problem mein",
        ],
        "es-en": [
            "Voy a la store para comprar some groceries",
            "El weather está muy nice hoy",
            "Tengo que hacer el homework antes de salir",
            "Mi familia va a tener una barbecue este weekend",
            "Necesito upgradar mi phone, el mío está muy old",
        ],
    }

    test_examples = examples.get(language_pair, examples["hi-en"])

    print("\n" + "=" * 70)
    print(f"TRANSLATION EXAMPLES ({language_pair.upper()})")
    print("=" * 70)

    for i, text in enumerate(test_examples, 1):
        start = time.time()
        translation = translator.translate(text, num_beams=4)
        elapsed = time.time() - start

        print(f"\n[Example {i}]")
        print(f"  Input:  {text}")
        print(f"  Output: {translation}")
        print(f"  Time:   {elapsed*1000:.0f}ms")

    print("\n" + "=" * 70)