app.py

"""
Translation interface using the MADLAD-400 3B model.
Translates English text to 22 production-ready languages from the MADLAD-400 paper.
"""

import csv
import tempfile
import warnings
from functools import lru_cache
from pathlib import Path

import gradio as gr
import spaces
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

from langmap.langid_mapping import langid_to_language

MODEL_NAME = "google/madlad400-3b-mt"


def _get_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    warnings.warn("No GPU available. Running on CPU — translation will be slow.", stacklevel=2)
    return torch.device("cpu")


@lru_cache(maxsize=1)
def _load_tokenizer() -> AutoTokenizer:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
    if tokenizer is None:
        raise RuntimeError(f"Failed to load tokenizer for {MODEL_NAME}")
    return tokenizer


@lru_cache(maxsize=1)
def _load_model() -> AutoModelForSeq2SeqLM:
    device = _get_device()
    dtype = torch.float16 if device.type == "cuda" else torch.float32
    return AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, dtype=dtype).to(device)


@lru_cache(maxsize=1)
def _build_language_mappings() -> tuple[dict[str, str], list[str]]:
    tokenizer = _load_tokenizer()
    vocab = tokenizer.get_vocab()
    name_to_code = {v: k for k, v in langid_to_language.items() if k in vocab}
    return name_to_code, sorted(name_to_code.keys())


@spaces.GPU
def translate(
    text: str,
    target_language_name: str,
    max_new_tokens: int = 512,
    num_beams: int = 1,
    temperature: float = 1.0,
) -> str:
    tokenizer = _load_tokenizer()
    model = _load_model()
    device = model.device

    name_to_code, _ = _build_language_mappings()
    target_code = name_to_code.get(target_language_name)
    if target_code is None:
        raise ValueError(f"Unsupported language: {target_language_name}")

    if num_beams > 1 and temperature != 1.0:
        gr.Info("Temperature has no effect when beam search is enabled (num_beams > 1).")

    input_ids = tokenizer(target_code + " " + text, return_tensors="pt").input_ids.to(device)

    generate_kwargs: dict = {"input_ids": input_ids, "max_new_tokens": max_new_tokens, "num_beams": num_beams}
    if num_beams == 1:
        generate_kwargs["do_sample"] = True
        generate_kwargs["temperature"] = temperature

    outputs = model.generate(**generate_kwargs)
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    if not isinstance(result, str):
        raise TypeError(f"Expected str from decode, got {type(result)}")
    return result


def _parse_input_file(file_path: str) -> tuple[list[str], list[dict] | None]:
    """Parse a .txt or .csv file into a list of texts to translate.

    Returns (texts, original_data). original_data is None for .txt files,
    or a list of row dicts for .csv files (used to preserve extra columns in output).
    """
    path = Path(file_path)
    ext = path.suffix.lower()

    if ext == ".txt":
        texts = [line for line in path.read_text().splitlines() if line.strip()]
        return texts, None

    if ext == ".csv":
        with open(path, newline="") as f:
            reader = csv.DictReader(f)
            if reader.fieldnames is None or "text" not in reader.fieldnames:
                raise ValueError("CSV file must have a 'text' column.")
            rows = list(reader)
        texts = [row["text"] for row in rows]
        return texts, rows

    raise ValueError(f"Unsupported file type: {ext}. Use .txt or .csv.")


def _write_output_file(translations: list[str], original_data: list[dict] | None, ext: str) -> str:
    """Write translation results to a temp file and return its path."""
    if ext == ".txt":
        f = tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False)
        f.write("\n".join(translations))
        f.close()
        return f.name

    # CSV: append translation column to original data
    f = tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False, newline="")
    assert original_data is not None
    fieldnames = list(original_data[0].keys()) + ["translation"]
    writer = csv.DictWriter(f, fieldnames=fieldnames)
    writer.writeheader()
    for row, translation in zip(original_data, translations):
        row["translation"] = translation
        writer.writerow(row)
    f.close()
    return f.name


def _build_table_data(texts: list[str], translations: list[str], original_data: list[dict] | None) -> dict:
    """Build a dict with headers and data for gr.Dataframe display."""
    if original_data is None:
        return {"headers": ["source", "translation"], "data": [[t, tr] for t, tr in zip(texts, translations)]}
    headers = list(original_data[0].keys()) + ["translation"]
    data = [[row[k] for k in original_data[0].keys()] + [tr] for row, tr in zip(original_data, translations)]
    return {"headers": headers, "data": data}


@spaces.GPU
def translate_batch(
    file,
    target_language_name: str,
    max_new_tokens: int = 512,
    num_beams: int = 1,
    temperature: float = 1.0,
    progress=gr.Progress(),
) -> tuple[dict, str]:
    if file is None:
        raise gr.Error("Please upload a file first.")

    file_path = file.name if hasattr(file, "name") else str(file)
    ext = Path(file_path).suffix.lower()
    texts, original_data = _parse_input_file(file_path)

    if not texts:
        raise gr.Error("File is empty or contains no text to translate.")

    if num_beams > 1 and temperature != 1.0:
        gr.Info("Temperature has no effect when beam search is enabled (num_beams > 1).")

    translations = []
    for i, text in enumerate(progress.tqdm(texts, desc="Translating")):
        result = translate(text, target_language_name, max_new_tokens, num_beams, temperature)
        translations.append(result)

    output_path = _write_output_file(translations, original_data, ext)
    table_data = _build_table_data(texts, translations, original_data)
    return table_data, output_path


def _build_demo() -> gr.Blocks:
    _, language_names = _build_language_mappings()

    with gr.Blocks(title="MADLAD-400 Translation") as demo:
        gr.Markdown(
            "# MADLAD-400 Translation\n"
            "Translate English into 22 production-ready languages using Google's MADLAD-400 3B model. "
            "[Paper](https://arxiv.org/pdf/2309.04662)"
        )

        with gr.Row():
            # --- Sidebar ---
            with gr.Column(scale=1, min_width=250):
                gr.Markdown("### Settings")

                target_language = gr.Dropdown(
                    choices=language_names,
                    value="French",
                    label="Target language",
                )
                max_new_tokens = gr.Slider(minimum=1, maximum=1024, value=512, step=1, label="Max new tokens")
                num_beams = gr.Slider(minimum=1, maximum=10, value=1, step=1, label="Num beams")
                temperature = gr.Slider(minimum=0.1, maximum=2.0, value=1.0, step=0.1, label="Temperature")

            # --- Main area ---
            with gr.Column(scale=3):
                with gr.Tab("Single"):
                    with gr.Row():
                        input_text = gr.Textbox(
                            label="English",
                            placeholder="Enter English text here",
                            lines=4,
                        )
                        output_text = gr.Textbox(
                            label="Translation",
                            lines=4,
                            buttons=["copy"],
                            interactive=False,
                        )

                    single_btn = gr.Button("Translate", variant="primary")

                    single_btn.click(
                        fn=translate,
                        inputs=[input_text, target_language, max_new_tokens, num_beams, temperature],
                        outputs=output_text,
                    )

                    gr.Examples(
                        examples=[
                            ["Hello, how are you today?", "French"],
                            ["The weather is beautiful.", "Spanish"],
                            ["Thank you very much.", "German"],
                            ["Where is the train station?", "Portuguese"],
                        ],
                        inputs=[input_text, target_language],
                    )

                with gr.Tab("Batch"):
                    with gr.Row():
                        batch_input = gr.File(
                            label="Upload file (.txt, one sentence per line, or .csv with 'text' column)",
                            file_types=[".txt", ".csv"],
                        )
                        batch_output = gr.File(label="Download translations", interactive=False)

                    batch_btn = gr.Button("Translate", variant="primary")

                    batch_preview = gr.Dataframe(label="Preview", interactive=False)

                    batch_btn.click(
                        fn=translate_batch,
                        inputs=[batch_input, target_language, max_new_tokens, num_beams, temperature],
                        outputs=[batch_preview, batch_output],
                    )

    return demo


demo = _build_demo()


def main() -> None:
    demo.launch(theme=gr.themes.Soft())


if __name__ == "__main__":
    main()