Update app.py

app.py

@@ -1,50 +1,20 @@
 import os
 import torch
 import spaces
-from collections.abc import Iterator
-from threading import Thread
 import gradio as gr
+from threading import Thread
+from collections.abc import Iterator
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 MAX_MAX_NEW_TOKENS = 4096
-DEFAULT_MAX_NEW_TOKENS = 2048
 MAX_INPUT_TOKEN_LENGTH = 4096
-
+DEFAULT_MAX_NEW_TOKENS = 2048
 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)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto", token=HF_TOKEN)
 tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B-Instruct")
 
-
 LANGUAGES = {
     "Hindi": "hin_Deva",
     "Bengali": "ben_Beng",
@@ -69,59 +39,33 @@ LANGUAGES = {
     "Bodo": "brx_Deva"
 }
 
-@spaces.GPU
-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 format_message_for_translation(message, target_lang):
+    return f"Translate the following text to {target_lang}: {message}"
 
 @spaces.GPU
-def generate(
-    tgt_lang: str,
+def translate_message(
     message: str,
+    chat_history: list[dict],
+    target_language: str = "Hindi",
     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}"})
+    
+    translation_request = format_message_for_translation(message, target_language)
+    print(f"Translation request: {translation_request}")
+    conversation.append({"role": "user", "content": translation_request})
 
     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.")
+        gr.Warning(f"Trimmed input 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)
+    streamer = TextIteratorStreamer(tokenizer, timeout=240.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         {"input_ids": input_ids},
         streamer=streamer,
@@ -154,84 +98,84 @@ def store_feedback(rating, feedback_text):
     return "Thank you for your feedback!"
 
 css = """
-      #col-container {max-width: 80%; margin-left: auto; margin-right: auto;}
-      #header {text-align: left;}
-      .message { font-size: 1.2em; }
-      #feedback-section { margin-top: 30px; border-top: 1px solid #ddd; padding-top: 20px; }
-      """
+body {
+    background-color: #f7f7f7;
+}
+.feedback-section {
+    margin-top: 30px;
+    border-top: 1px solid #ddd;
+    padding-top: 20px;
+}
+.container {
+    max-width: 90%;
+    margin: 0 auto;
+}
+.language-selector {
+    margin-bottom: 20px;
+    padding: 10px;
+    background-color: #ffffff;
+    border-radius: 8px;
+    box-shadow: 0 2px 5px rgba(0,0,0,0.1);
+}
+.advanced-options {
+    margin-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"
-                )
+DESCRIPTION = """\
+IndicTrans3 is the latest state-of-the-art (SOTA) translation model from AI4Bharat, designed to handle translations across <b>22 Indic languages</b> with high accuracy. It supports <b>document-level machine translation (MT)</b> and is built to match the performance of other leading SOTA models. <br>
+📢 <b>Training data will be released soon!</b>  
+<h3>🔹 Features</h3>  
+✅ Supports <b>22 Indic languages</b>  
+✅ Enables <b>document-level translation</b>  
+✅ Achieves <b>SOTA performance</b> in Indic MT  
+✅ Optimized for <b>real-world applications</b>  
+<h3>🚀 Try It Out!</h3>  
+1️⃣ Enter text in any supported language  
+2️⃣ Select the target language  
+3️⃣ Click <b>Translate</b> and get high-quality results!  
+Built for <b>linguistic diversity and accessibility</b>, IndicTrans3 is a major step forward in <b>Indic language AI</b>.  
+💡 <b>Source:</b> AI4Bharat | Powered by Hugging Face  
+"""
+
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_classes="container"):
+        gr.Markdown("# 🌏 IndicTrans3-beta 🚀: Multilingual Translation for 22 Indic Languages </center>")
+        gr.Markdown(DESCRIPTION)
         
-        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
+        target_language = gr.Dropdown(
+            list(LANGUAGES.keys()), 
+            value="Hindi",
+            label="Which language would you like to translate to?",
+            elem_id="language-dropdown"
         )
+    
+        chatbot = gr.Chatbot(height=400, elem_id="chatbot")
+        
+        with gr.Row():
+            msg = gr.Textbox(
+                placeholder="Enter text to translate...",
+                show_label=False,
+                container=False,
+                scale=9
+            )
+            submit_btn = gr.Button("Translate", scale=1)
         
         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."]
+                "The Taj Mahal stands majestically along the banks of river Yamuna, a timeless symbol of eternal love.",
+                "Kumbh Mela is the world's largest gathering of people, where millions of pilgrims bathe in sacred rivers for spiritual purification.",
+                "India's classical dance forms like Bharatanatyam, Kathak, and Odissi beautifully blend rhythm, expression, and storytelling.",
+                "Ayurveda, the ancient Indian medical system, focuses on holistic wellness through natural herbs and balanced living.",
+                "During Diwali, homes across India are decorated with oil lamps, colorful rangoli patterns, and twinkling lights to celebrate the victory of light over darkness."
             ],
-            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
+            inputs=msg
         )
         
-        with gr.Column(elem_id="feedback-section"):
+        
+        with gr.Accordion("Provide Feedback", open=True):
             gr.Markdown("## Rate Translation & Provide Feedback 📝")
-            gr.Markdown("Help us improve the translation quality by providing your feedback and rating.")
+            gr.Markdown("Help us improve the translation quality by providing your feedback.")
             with gr.Row():
                 rating = gr.Radio(
                     ["1", "2", "3", "4", "5"],
@@ -246,11 +190,93 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
             
             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
+
+        with gr.Accordion("Advanced Options", open=False, elem_classes="advanced-options"):
+            max_new_tokens = gr.Slider(
+                label="Max new tokens",
+                minimum=1,
+                maximum=MAX_MAX_NEW_TOKENS,
+                step=1,
+                value=DEFAULT_MAX_NEW_TOKENS,
+            )
+            temperature = gr.Slider(
+                label="Temperature",
+                minimum=0.1,
+                maximum=1.0,
+                step=0.1,
+                value=0.1,
+            )
+            top_p = gr.Slider(
+                label="Top-p (nucleus sampling)",
+                minimum=0.05,
+                maximum=1.0,
+                step=0.05,
+                value=0.9,
+            )
+            top_k = gr.Slider(
+                label="Top-k",
+                minimum=1,
+                maximum=100,
+                step=1,
+                value=50,
+            )
+            repetition_penalty = gr.Slider(
+                label="Repetition penalty",
+                minimum=1.0,
+                maximum=2.0,
+                step=0.05,
+                value=1.0,
             )
+        
+        chat_state = gr.State([])
+        
+        def user(user_message, history, target_lang):
+            return "", history + [[user_message, None]]
+        
+        def bot(history, target_lang, max_tokens, temp, top_p_val, top_k_val, rep_penalty):
+            user_message = history[-1][0]
+            history[-1][1] = ""
+            
+            for chunk in translate_message(
+                user_message, 
+                history[:-1], 
+                target_lang, 
+                max_tokens, 
+                temp, 
+                top_p_val, 
+                top_k_val, 
+                rep_penalty
+            ):
+                history[-1][1] = chunk
+                yield history
+        
+        msg.submit(
+            user,
+            [msg, chatbot, target_language],
+            [msg, chatbot],
+            queue=False
+        ).then(
+            bot,
+            [chatbot, target_language, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
+            chatbot
+        )
+        
+        submit_btn.click(
+            user,
+            [msg, chatbot, target_language],
+            [msg, chatbot],
+            queue=False
+        ).then(
+            bot,
+            [chatbot, target_language, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
+            chatbot
+        )
+        
+        feedback_submit.click(
+            fn=store_feedback, 
+            inputs=[rating, feedback_text], 
+            outputs=feedback_result
+        )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()