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