app.py

#### Former version using Helsinki NLP

import gradio as gr
import pandas as pd
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import io
import re
from typing import Tuple, Optional, List

# NLLB-200 language codes (BCP-47 + script tags)
LANGUAGE_PAIRS = {
    "English": "eng_Latn",
    "Spanish": "spa_Latn",
    "French": "fra_Latn",
    "German": "deu_Latn",
    "Italian": "ita_Latn",
    "Portuguese": "por_Latn",
    "Russian": "rus_Cyrl",
    "Chinese (Simplified)": "zho_Hans",
    "Chinese (Traditional)": "zho_Hant",
    "Japanese": "jpn_Jpan",
    "Korean": "kor_Hang",
    "Arabic": "arb_Arab",
    "Hindi": "hin_Deva",
    "Dutch": "nld_Latn",
    "Polish": "pol_Latn",
    "Turkish": "tur_Latn",
    "Vietnamese": "vie_Latn",
    "Indonesian": "ind_Latn",
    "Thai": "tha_Thai",
    "Swedish": "swe_Latn",
    "Czech": "ces_Latn",
    "Romanian": "ron_Latn",
    "Hungarian": "hun_Latn",
    "Ukrainian": "ukr_Cyrl",
    "Greek": "ell_Grek",
    "Hebrew": "heb_Hebr",
    "Bengali": "ben_Beng",
    "Swahili": "swh_Latn",
    "Finnish": "fin_Latn",
    "Norwegian": "nob_Latn",
}

# Model options
MODEL_OPTIONS = {
    "NLLB-200 Distilled 600M (Recommended)": "facebook/nllb-200-distilled-600M",
    "NLLB-200 Distilled 1.3B (Higher Quality)": "facebook/nllb-200-distilled-1.3B",
    "NLLB-200 1.3B": "facebook/nllb-200-1.3B",
}

# Cache for loaded pipelines
translation_cache = {}


def segment_text(text: str) -> List[str]:
    """Segment text into sentences for translation."""
    segments = re.split(r'([.!?]+\s+|[.!?]+$|\n+)', text)

    result = []
    i = 0
    while i < len(segments):
        if segments[i].strip():
            segment = segments[i]
            if i + 1 < len(segments) and re.match(r'^[.!?]+\s*$|^\n+$', segments[i + 1]):
                segment += segments[i + 1]
                i += 2
            else:
                i += 1
            result.append(segment.strip())
        else:
            i += 1

    return result if result else [text]


def load_translator(model_name: str):
    """Load or retrieve a cached NLLB translation pipeline."""
    if model_name not in translation_cache:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        translation_cache[model_name] = {"tokenizer": tokenizer, "model": model}
    return translation_cache[model_name]


def translate_text(text: str, tokenizer, model, src_lang: str, tgt_lang: str) -> str:
    """Translate a single piece of text using the NLLB model."""
    tokenizer.src_lang = src_lang
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    forced_bos_token_id = tokenizer.convert_tokens_to_ids(tgt_lang)
    outputs = model.generate(
        **inputs,
        forced_bos_token_id=forced_bos_token_id,
        max_new_tokens=512,
        num_beams=4,
        early_stopping=True,
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)


def translate_segments(segments: List[str], tokenizer, model, src_lang: str, tgt_lang: str) -> str:
    """Translate each segment and join them."""
    translated = []
    for seg in segments:
        if not seg.strip():
            translated.append(seg)
            continue
        try:
            translated.append(translate_text(seg, tokenizer, model, src_lang, tgt_lang))
        except Exception as e:
            translated.append(f"[Error: {seg}]")
    return " ".join(translated)


def read_file(file) -> Tuple[Optional[pd.DataFrame], Optional[str], Optional[list]]:
    """Read uploaded CSV or Excel file."""
    if file is None:
        return None, "Please upload a file", None
    try:
        ext = file.name.split(".")[-1].lower()
        if ext == "csv":
            df = pd.read_csv(file.name)
        elif ext in ["xlsx", "xls"]:
            df = pd.read_excel(file.name)
        else:
            return None, "Unsupported format. Upload CSV or Excel.", None

        if df.empty:
            return None, "The uploaded file is empty", None

        return df, None, df.columns.tolist()
    except Exception as e:
        return None, f"Error reading file: {str(e)}", None


def translate_column(
    file,
    column_name: str,
    source_lang: str,
    target_lang: str,
    model_choice: str,
    output_format: str,
    progress=gr.Progress(),
) -> Tuple[Optional[str], Optional[str], Optional[pd.DataFrame]]:
    """Main translation function."""
    if file is None:
        return None, "Please upload a file", None
    if not column_name or column_name == "Select a column":
        return None, "Please select a column to translate", None
    if source_lang == target_lang:
        return None, "Source and target languages must be different", None

    df, error, _ = read_file(file)
    if error:
        return None, error, None
    if column_name not in df.columns:
        return None, f"Column '{column_name}' not found in file", None

    src_code = LANGUAGE_PAIRS[source_lang]
    tgt_code = LANGUAGE_PAIRS[target_lang]
    model_name = MODEL_OPTIONS[model_choice]

    progress(0, desc=f"Loading {model_choice}...")
    try:
        translator_dict = load_translator(model_name)
        tokenizer = translator_dict["tokenizer"]
        model = translator_dict["model"]
    except Exception as e:
        return None, f"Failed to load model: {str(e)}", None

    progress(0.1, desc="Starting translation...")
    translated_texts = []
    total_rows = len(df)

    for idx, text in enumerate(df[column_name]):
        progress((idx + 1) / total_rows, desc=f"Translating row {idx + 1}/{total_rows}")
        if pd.isna(text) or str(text).strip() == "":
            translated_texts.append("")
        else:
            try:
                segments = segment_text(str(text).strip())
                translated = translate_segments(segments, tokenizer, model, src_code, tgt_code)
                translated_texts.append(translated)
            except Exception as e:
                translated_texts.append(f"[Translation Error: {str(e)}]")

    new_col = f"{column_name}_{tgt_code}"
    df[new_col] = translated_texts

    progress(1.0, desc="Generating output file...")
    output = io.BytesIO()
    if output_format == "CSV":
        df.to_csv(output, index=False)
        file_name = f"translated_{src_code}_to_{tgt_code}.csv"
    else:
        df.to_excel(output, index=False, engine="openpyxl")
        file_name = f"translated_{src_code}_to_{tgt_code}.xlsx"

    output.seek(0)
    with open(file_name, "wb") as f:
        f.write(output.getvalue())

    msg = f"✅ Successfully translated {total_rows} rows from {source_lang} to {target_lang} using {model_choice}"
    return file_name, msg, df.head(10)


def update_column_choices(file):
    """Update column dropdown when file is uploaded."""
    if file is None:
        return gr.Dropdown(choices=["Select a column"], value="Select a column")
    df, error, columns = read_file(file)
    if error or not columns:
        return gr.Dropdown(choices=["Select a column"], value="Select a column")
    return gr.Dropdown(choices=columns, value=columns[0])


# ── UI ──────────────────────────────────────────────────────────────────────
with gr.Blocks(title="Spreadsheet Column Translator", theme=gr.themes.Soft()) as app:
    gr.Markdown(
        """
        # 🌍 Spreadsheet Column Translator
        Upload a CSV or Excel file, pick a column, and get high-quality translations powered by
        **Meta's NLLB-200** — supporting 200+ languages with state-of-the-art accuracy.
        """
    )

    with gr.Row():
        with gr.Column():
            file_input = gr.File(
                label="📁 Upload Spreadsheet (CSV or Excel)",
                file_types=[".csv", ".xlsx", ".xls"],
            )
            column_dropdown = gr.Dropdown(
                choices=["Select a column"],
                value="Select a column",
                label="📊 Select Column to Translate",
                interactive=True,
            )
            with gr.Row():
                source_lang = gr.Dropdown(
                    choices=list(LANGUAGE_PAIRS.keys()),
                    value="English",
                    label="🗣️ Source Language",
                )
                target_lang = gr.Dropdown(
                    choices=list(LANGUAGE_PAIRS.keys()),
                    value="Spanish",
                    label="🎯 Target Language",
                )
            model_choice = gr.Dropdown(
                choices=list(MODEL_OPTIONS.keys()),
                value="NLLB-200 Distilled 600M (Recommended)",
                label="🤖 Translation Model",
            )
            output_format = gr.Radio(
                choices=["CSV", "Excel"], value="CSV", label="💾 Output Format"
            )
            translate_btn = gr.Button("🔄 Translate", variant="primary", size="lg")

        with gr.Column():
            status_output = gr.Textbox(label="📝 Status", lines=3, interactive=False)
            file_output = gr.File(label="⬇️ Download Translated File")

    gr.Markdown("### 👀 Preview (First 10 rows)")
    preview_output = gr.Dataframe(label="Translated Data Preview", wrap=True)

    gr.Markdown(
        """
        ---
        **Models:**
        - **NLLB-200 Distilled 600M** — Fast, great quality. Best for most use cases.
        - **NLLB-200 Distilled 1.3B** — Slower but higher quality, especially for rare languages.
        - **NLLB-200 1.3B** — Full (non-distilled) 1.3B model for maximum accuracy.

        **Tip:** The first run for each model downloads it from Hugging Face — subsequent runs use the cache.
        """
    )

    file_input.change(fn=update_column_choices, inputs=[file_input], outputs=[column_dropdown])
    translate_btn.click(
        fn=translate_column,
        inputs=[file_input, column_dropdown, source_lang, target_lang, model_choice, output_format],
        outputs=[file_output, status_output, preview_output],
    )

if __name__ == "__main__":
    app.launch()