Spaces:
Sleeping
Sleeping
| 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() |