Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| import torch | |
| import spaces | |
| from collections.abc import Iterator | |
| from threading import Thread | |
| import gradio as gr | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer | |
| MAX_MAX_NEW_TOKENS = 4096 | |
| DEFAULT_MAX_NEW_TOKENS = 2048 | |
| MAX_INPUT_TOKEN_LENGTH = 4096 | |
| HF_TOKEN = os.environ['HF_TOKEN'] | |
| DESCRIPTION = """\ | |
| ## 🌏 IndicTrans3-beta 🚀: Multilingual Translation for 22 Indic Languages | |
| IndicTrans3 is the latest state-of-the-art (SOTA) translation model from AI4Bharat, designed to handle translations across **22 Indic languages** with high accuracy. It supports **document-level machine translation (MT)** and is built to match the performance of other leading SOTA models. | |
| 📢 **Training data will be released soon!** | |
| ### 🔹 Features | |
| ✅ Supports **22 Indic languages** | |
| ✅ Enables **document-level translation** | |
| ✅ Achieves **SOTA performance** in Indic MT | |
| ✅ Optimized for **real-world applications** | |
| ### 🚀 Try It Out! | |
| 1️⃣ Enter text in any supported language | |
| 2️⃣ Select the target language | |
| 3️⃣ Click **Translate** and get high-quality results! | |
| Built for **linguistic diversity and accessibility**, IndicTrans3 is a major step forward in **Indic language AI**. | |
| 💡 **Source:** AI4Bharat | Powered by Hugging Face | |
| """ | |
| # if not torch.cuda.is_available(): | |
| # DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>" | |
| # if torch.cuda.is_available(): | |
| model_id = "ai4bharat/IndicTrans3-beta" | |
| model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto", offload_folder="offload", token=HF_TOKEN) | |
| tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B-Instruct") | |
| LANGUAGES = { | |
| "Hindi": "hin_Deva", | |
| "Bengali": "ben_Beng", | |
| "Telugu": "tel_Telu", | |
| "Marathi": "mar_Deva", | |
| "Tamil": "tam_Taml", | |
| "Urdu": "urd_Arab", | |
| "Gujarati": "guj_Gujr", | |
| "Kannada": "kan_Knda", | |
| "Odia": "ori_Orya", | |
| "Malayalam": "mal_Mlym", | |
| "Punjabi": "pan_Guru", | |
| "Assamese": "asm_Beng", | |
| "Maithili": "mai_Mith", | |
| "Santali": "sat_Olck", | |
| "Kashmiri": "kas_Arab", | |
| "Nepali": "nep_Deva", | |
| "Sindhi": "snd_Arab", | |
| "Konkani": "kok_Deva", | |
| "Dogri": "dgo_Deva", | |
| "Manipuri": "mni_Beng", | |
| "Bodo": "brx_Deva" | |
| } | |
| def generate_for_examples( | |
| tgt_lang: str, | |
| message: str, | |
| max_new_tokens: int = 1024, | |
| temperature: float = 0.6, | |
| top_p: float = 0.9, | |
| top_k: int = 50, | |
| repetition_penalty: float = 1.2, | |
| ) -> str: | |
| conversation = [] | |
| conversation.append({"role": "user", "content": f"Translate the following text to {tgt_lang}: {message}"}) | |
| input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt", add_generation_prompt=True) | |
| if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH: | |
| input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:] | |
| input_ids = input_ids.to(model.device) | |
| outputs = model.generate( | |
| input_ids=input_ids, | |
| max_new_tokens=max_new_tokens, | |
| do_sample=True, | |
| top_p=top_p, | |
| top_k=top_k, | |
| temperature=temperature, | |
| num_beams=1, | |
| repetition_penalty=repetition_penalty, | |
| ) | |
| return tokenizer.decode(outputs[0][input_ids.shape[1]:], skip_special_tokens=True) | |
| def generate( | |
| tgt_lang: str, | |
| message: str, | |
| max_new_tokens: int = 1024, | |
| temperature: float = 0.6, | |
| top_p: float = 0.9, | |
| top_k: int = 50, | |
| repetition_penalty: float = 1.2, | |
| ) -> Iterator[str]: | |
| conversation = [] | |
| conversation.append({"role": "user", "content": f"Translate the following text to {tgt_lang}: {message}"}) | |
| input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt", add_generation_prompt=True) | |
| if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH: | |
| input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:] | |
| gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.") | |
| input_ids = input_ids.to(model.device) | |
| streamer = TextIteratorStreamer(tokenizer, timeout=180.0, skip_prompt=True, skip_special_tokens=True) | |
| generate_kwargs = dict( | |
| {"input_ids": input_ids}, | |
| streamer=streamer, | |
| max_new_tokens=max_new_tokens, | |
| do_sample=True, | |
| top_p=top_p, | |
| top_k=top_k, | |
| temperature=temperature, | |
| num_beams=1, | |
| repetition_penalty=repetition_penalty, | |
| ) | |
| t = Thread(target=model.generate, kwargs=generate_kwargs) | |
| t.start() | |
| outputs = [] | |
| for text in streamer: | |
| outputs.append(text) | |
| yield "".join(outputs) | |
| def store_feedback(rating, feedback_text): | |
| if not rating: | |
| gr.Warning("Please select a rating before submitting feedback.", duration=5) | |
| return None | |
| if not feedback_text or feedback_text.strip() == "": | |
| gr.Warning("Please provide some feedback before submitting.", duration=5) | |
| return None | |
| gr.Info("Feedback submitted successfully!") | |
| return "Thank you for your feedback!" | |
| css = """ | |
| #col-container {max-width: 80%; margin-left: auto; margin-right: auto;} | |
| #header {text-align: center;} | |
| .message { font-size: 1.2em; } | |
| #feedback-section { margin-top: 30px; border-top: 1px solid #ddd; padding-top: 20px; } | |
| """ | |
| with gr.Blocks(theme=gr.themes.Default(), css=css) as demo: | |
| gr.Markdown(DESCRIPTION, elem_id="header") | |
| gr.Markdown("Translate text between multiple Indic languages using the latest IndicTrans3 model from AI4Bharat. This model is trained on the --- dataset and supports translation to 22 Indic languages. Setting a state-of-the-art benchmark on multiple translation tasks, IndicTrans3 is a powerful model that can handle complex translation tasks with ease.", elem_id="description") | |
| with gr.Column(elem_id="col-container"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| text_input = gr.Textbox( | |
| placeholder="Enter text to translate...", | |
| label="Input text", | |
| lines=10, | |
| max_lines=100, | |
| elem_id="input-text" | |
| ) | |
| with gr.Column(): | |
| tgt_lang = gr.Dropdown( | |
| list(LANGUAGES.keys()), | |
| value="Hindi", | |
| label="Translate To", | |
| elem_id="translate-to" | |
| ) | |
| text_output = gr.Textbox( | |
| label="", | |
| lines=10, | |
| max_lines=100, | |
| elem_id="output-text" | |
| ) | |
| btn_submit = gr.Button("Translate") | |
| btn_submit.click( | |
| fn=generate, | |
| inputs=[ | |
| tgt_lang, | |
| text_input, | |
| gr.Number(value=4096, visible=False), | |
| gr.Number(value=0.1, visible=False), | |
| gr.Number(value=0.9, visible=False), | |
| gr.Number(value=50, visible=False), | |
| gr.Number(value=1.0, visible=False) | |
| ], | |
| outputs=text_output | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["Telugu", "Hello, how are you today? I hope you're doing well."], | |
| ["Punjabi", "Hello, how are you today? I hope you're doing well."], | |
| ["Hindi", "Hello, how are you today? I hope you're doing well."], | |
| ["Marathi", "Hello, how are you today? I hope you're doing well."], | |
| ["Malayalam", "Hello, how are you today? I hope you're doing well."] | |
| ], | |
| inputs=[ | |
| tgt_lang, | |
| text_input, | |
| gr.Number(value=4096, visible=False), | |
| gr.Number(value=0.1, visible=False), | |
| gr.Number(value=0.9, visible=False), | |
| gr.Number(value=50, visible=False), | |
| gr.Number(value=1.0, visible=False) | |
| ], | |
| outputs=text_output, | |
| fn=generate_for_examples, | |
| cache_examples=True, | |
| examples_per_page=5 | |
| ) | |
| with gr.Column(elem_id="feedback-section"): | |
| gr.Markdown("## Rate Translation & Provide Feedback 📝") | |
| gr.Markdown("Help us improve the translation quality by providing your feedback and rating.") | |
| with gr.Row(): | |
| rating = gr.Radio( | |
| ["1", "2", "3", "4", "5"], | |
| label="Translation Rating (1-5)" | |
| ) | |
| feedback_text = gr.Textbox( | |
| placeholder="Share your feedback about the translation...", | |
| label="Feedback", | |
| lines=3 | |
| ) | |
| feedback_submit = gr.Button("Submit Feedback") | |
| feedback_result = gr.Textbox(label="", visible=False) | |
| feedback_submit.click( | |
| fn=store_feedback, | |
| inputs=[rating, feedback_text], | |
| outputs=feedback_result | |
| ) | |
| demo.launch() |