app.py

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

# ----------------------------
# Models
# ----------------------------
# Translation model (100+ languages)
TRANS_MODEL_NAME = "facebook/m2m100_418M"

# Summarization model (English)
SUMM_MODEL_NAME = "sshleifer/distilbart-cnn-12-6"

_translation_model = None
_translation_tokenizer = None
_summarizer = None

# Curated list of widely used languages supported by M2M100.
# Keys are human-friendly names; values are M2M100 language codes.
LANGUAGES = {
    "English": "en",
    "French": "fr",
    "German": "de",
    "Spanish": "es",
    "Portuguese": "pt",
    "Italian": "it",
    "Dutch": "nl",
    "Swedish": "sv",
    "Danish": "da",
    "Norwegian": "no",
    "Finnish": "fi",
    "Polish": "pl",
    "Czech": "cs",
    "Slovak": "sk",
    "Hungarian": "hu",
    "Romanian": "ro",
    "Bulgarian": "bg",
    "Greek": "el",
    "Turkish": "tr",
    "Russian": "ru",
    "Ukrainian": "uk",
    "Serbian": "sr",
    "Croatian": "hr",
    "Slovenian": "sl",
    "Albanian": "sq",
    "Macedonian": "mk",
    "Bosnian": "bs",
    "Arabic": "ar",
    "Hebrew": "he",
    "Persian (Farsi)": "fa",
    "Urdu": "ur",
    "Hindi": "hi",
    "Bengali": "bn",
    "Punjabi": "pa",
    "Gujarati": "gu",
    "Marathi": "mr",
    "Tamil": "ta",
    "Telugu": "te",
    "Kannada": "kn",
    "Malayalam": "ml",
    "Sinhala": "si",
    "Nepali": "ne",
    "Indonesian": "id",
    "Malay": "ms",
    "Vietnamese": "vi",
    "Thai": "th",
    "Filipino (Tagalog)": "tl",
    "Chinese (Simplified)": "zh",
    "Japanese": "ja",
    "Korean": "ko",
    "Mongolian": "mn",
    "Kazakh": "kk",
    "Uzbek": "uz",
    "Azerbaijani": "az",
    "Georgian": "ka",
    "Armenian": "hy",
    "Amharic": "am",
    "Swahili": "sw",
    "Hausa": "ha",
    "Yoruba": "yo",
    "Igbo": "ig",
    "Zulu": "zu",
    "Xhosa": "xh",
    "Afrikaans": "af",
    "Irish": "ga",
    "Welsh": "cy",
    "Basque": "eu",
    "Catalan": "ca",
    "Galician": "gl",
    "Estonian": "et",
    "Latvian": "lv",
    "Lithuanian": "lt",
}

LANG_NAMES = sorted(LANGUAGES.keys())
DEFAULT_SRC = "English"
DEFAULT_TGT = "French"

def load_translation_model():
    global _translation_model, _translation_tokenizer
    if _translation_model is None or _translation_tokenizer is None:
        _translation_tokenizer = AutoTokenizer.from_pretrained(TRANS_MODEL_NAME)
        _translation_model = AutoModelForSeq2SeqLM.from_pretrained(TRANS_MODEL_NAME)
        _translation_model.eval()
    return _translation_model, _translation_tokenizer

def load_summarizer():
    global _summarizer
    if _summarizer is None:
        _summarizer = pipeline(
            "summarization",
            model=SUMM_MODEL_NAME,
        )
    return _summarizer

# ----------------------------
# Inference functions
# ----------------------------
def translate(text: str, src_lang_name: str, tgt_lang_name: str, max_new_tokens: int = 256, num_beams: int = 4):
    if not text or not text.strip():
        return ""
    if src_lang_name == tgt_lang_name:
        return text

    try:
        model, tokenizer = load_translation_model()
        src_code = LANGUAGES[src_lang_name]
        tgt_code = LANGUAGES[tgt_lang_name]

        tokenizer.src_lang = src_code
        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=1024)
        with torch.no_grad():
            generated_tokens = model.generate(
                **inputs,
                forced_bos_token_id=tokenizer.get_lang_id(tgt_code),
                max_new_tokens=max_new_tokens,
                num_beams=num_beams,
                length_penalty=1.0,
                early_stopping=True,
            )
        out = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
        return out[0] if out else ""
    except Exception as e:
        return f"Error during translation: {e}"

def summarize(text: str, min_length: int = 30, max_length: int = 130):
    if not text or not text.strip():
        return ""
    if max_length <= min_length:
        return "max_length must be greater than min_length."
    try:
        summarizer = load_summarizer()
        # The summarization pipeline will handle necessary truncation.
        output = summarizer(text, min_length=min_length, max_length=max_length, do_sample=False)
        return output[0]["summary_text"] if output else ""
    except Exception as e:
        return f"Error during summarization: {e}"

# ----------------------------
# Gradio UI
# ----------------------------
TITLE = "Multipurpose NLP: Translation and Summarization"
DESCRIPTION = """
This Space provides:
- Multi-language translation using M2M100 (100+ languages)
- Abstractive summarization in English using DistilBART CNN

Notes:
- Translation: pick source and target languages. For very long inputs, the model may truncate.
- Summarization: best for English. For long articles, try increasing max_length.
"""

with gr.Blocks(title=TITLE, theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"# {TITLE}")
    gr.Markdown(DESCRIPTION)

    with gr.Tabs():
        with gr.Tab("Translate"):
            with gr.Row():
                src = gr.Dropdown(choices=LANG_NAMES, value=DEFAULT_SRC, label="Source language")
                tgt = gr.Dropdown(choices=LANG_NAMES, value=DEFAULT_TGT, label="Target language")
            with gr.Row():
                inp = gr.Textbox(label="Input text", lines=8, placeholder="Enter text to translate...")
            with gr.Accordion("Advanced options", open=False):
                max_new = gr.Slider(32, 512, value=256, step=8, label="Max new tokens")
                beams = gr.Slider(1, 8, value=4, step=1, label="Beam width (higher = better but slower)")
            with gr.Row():
                btn_swap = gr.Button("Swap Languages ↔️", variant="secondary")
                btn_translate = gr.Button("Translate", variant="primary")
            out = gr.Textbox(label="Translated text", lines=8)

            def _swap(a, b):
                return b, a
            btn_swap.click(_swap, inputs=[src, tgt], outputs=[src, tgt])
            btn_translate.click(translate, inputs=[inp, src, tgt, max_new, beams], outputs=[out])

            gr.Examples(
                examples=[
                    ["Hello, how are you?", "English", "Spanish"],
                    ["Este libro es muy interesante.", "Spanish", "English"],
                    ["Je cherche un bon restaurant près d'ici.", "French", "English"],
                    ["机器学习正在改变世界。", "Chinese (Simplified)", "English"],
                ],
                inputs=[inp, src, tgt],
                label="Try examples",
            )

        with gr.Tab("Summarize"):
            s_in = gr.Textbox(label="Input text (English)", lines=12, placeholder="Paste article or long text...")
            with gr.Row():
                s_min = gr.Slider(10, 120, value=30, step=2, label="Min summary length")
                s_max = gr.Slider(40, 300, value=130, step=2, label="Max summary length")
            s_btn = gr.Button("Summarize", variant="primary")
            s_out = gr.Textbox(label="Summary", lines=10)

            s_btn.click(summarize, inputs=[s_in, s_min, s_max], outputs=[s_out])

            gr.Examples(
                examples=[
                    ["The Apollo program was the third United States human spaceflight program carried out by NASA, which succeeded in preparing and landing the first humans on the Moon from 1969 to 1972. It was first conceived during Dwight D. Eisenhower's administration as a three-person spacecraft to follow the one-person Project Mercury..."],
                ],
                inputs=[s_in],
                label="Example input",
            )

    gr.Markdown(
        "Models: facebook/m2m100_418M (translation), sshleifer/distilbart-cnn-12-6 (summarization). "
        "Running on CPU may be slow for very long texts."
    )

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