update app

app.py

@@ -1,6 +1,7 @@
 import os
 import torch
 import spaces
+import psycopg2
 import gradio as gr
 from threading import Thread
 from collections.abc import Iterator
@@ -9,39 +10,126 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
 MAX_MAX_NEW_TOKENS = 4096
 MAX_INPUT_TOKEN_LENGTH = 4096
 DEFAULT_MAX_NEW_TOKENS = 2048
-HF_TOKEN = os.environ['HF_TOKEN']
+HF_TOKEN = os.environ["HF_TOKEN"]
 
 model_id = "ai4bharat/IndicTrans3-beta"
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto", 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",
-    "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"
-}
+
+LANGUAGES = [
+    "Hindi",
+    "Bengali",
+    "Telugu",
+    "Marathi",
+    "Tamil",
+    "Urdu",
+    "Gujarati",
+    "Kannada",
+    "Odia",
+    "Malayalam",
+    "Punjabi",
+    "Assamese",
+    "Maithili",
+    "Santali",
+    "Kashmiri",
+    "Nepali",
+    "Sindhi",
+    "Konkani",
+    "Dogri",
+    "Manipuri",
+    "Bodo",
+]
+
 
 def format_message_for_translation(message, target_lang):
     return f"Translate the following text to {target_lang}: {message}"
 
+
+def store_feedback(rating, feedback_text, chat_history, tgt_lang):
+    try:
+
+        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
+
+        if not chat_history:
+            gr.Warning(
+                "Please provide the input text before submitting feedback.", duration=5
+            )
+            return None
+
+        if len(chat_history[0]) < 2:
+            gr.Warning(
+                "Please translate the input text before submitting feedback.",
+                duration=5,
+            )
+            return None
+
+        conn = psycopg2.connect(
+            host=os.getenv("DB_HOST"),
+            database=os.getenv("DB_NAME"),
+            user=os.getenv("DB_USER"),
+            password=os.getenv("DB_PASSWORD"),
+            port=os.getenv("DB_PORT"),
+        )
+
+        cursor = conn.cursor()
+
+        insert_query = """
+        INSERT INTO feedback 
+        (tgt_lang, rating, feedback_txt, chat_history)
+        VALUES (%s, %s, %s, %s)
+        """
+
+        cursor.execute(
+            insert_query, (tgt_lang, int(rating), feedback_text, chat_history)
+        )
+
+        conn.commit()
+
+        cursor.close()
+        conn.close()
+
+        gr.Info("Thank you for your feedback! 🙏", duration=5)
+
+    except:
+        gr.Error(
+            "An error occurred while storing feedback. Please try again later.",
+            duration=5,
+        )
+
+
+def store_output(tgt_lang, input_text, output_text):
+
+    conn = psycopg2.connect(
+        host=os.getenv("DB_HOST"),
+        database=os.getenv("DB_NAME"),
+        user=os.getenv("DB_USER"),
+        password=os.getenv("DB_PASSWORD"),
+        port=os.getenv("DB_PORT"),
+    )
+
+    cursor = conn.cursor()
+
+    insert_query = """
+    INSERT INTO translation
+    (input_txt, output_txt, tgt_lang)
+    VALUES (%s, %s, %s)
+    """
+
+    cursor.execute(insert_query, (input_text, output_text, tgt_lang))
+
+    conn.commit()
+    cursor.close()
+
+
 @spaces.GPU
 def translate_message(
     message: str,
@@ -54,18 +142,24 @@ def translate_message(
     repetition_penalty: float = 1.2,
 ) -> Iterator[str]:
     conversation = []
-    
+
     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)
+    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 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=240.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,
@@ -85,17 +179,8 @@ def translate_message(
         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!"
+    store_output(target_language, message, "".join(outputs))
+
 
 css = """
 # body {
@@ -140,58 +225,62 @@ Built for <b>linguistic diversity and accessibility</b>, IndicTrans3 is a major
 
 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(
+            "# 🌏 IndicTrans3-beta 🚀: Multilingual Translation for 22 Indic Languages </center>"
+        )
         gr.Markdown(DESCRIPTION)
-        
+
         target_language = gr.Dropdown(
-            list(LANGUAGES.keys()), 
+            LANGUAGES,
             value="Hindi",
             label="Which language would you like to translate to?",
-            elem_id="language-dropdown"
+            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
+                scale=9,
             )
             submit_btn = gr.Button("Translate", scale=1)
-        
+
         gr.Examples(
             examples=[
                 "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."
+                "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=msg
+            inputs=msg,
         )
-        
-        
+
         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.")
+            gr.Markdown(
+                "Help us improve the translation quality by providing your feedback."
+            )
             with gr.Row():
                 rating = gr.Radio(
-                    ["1", "2", "3", "4", "5"],
-                    label="Translation Rating (1-5)"
+                    ["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
+                lines=3,
             )
-            
+
             feedback_submit = gr.Button("Submit Feedback")
             feedback_result = gr.Textbox(label="", visible=False)
 
-        with gr.Accordion("Advanced Options", open=False, elem_classes="advanced-options"):
+        with gr.Accordion(
+            "Advanced Options", open=False, elem_classes="advanced-options"
+        ):
             max_new_tokens = gr.Slider(
                 label="Max new tokens",
                 minimum=1,
@@ -227,56 +316,66 @@ with gr.Blocks(css=css) as demo:
                 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):
+
+        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
+                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
+            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
+            [
+                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
+            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
+            [
+                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
+            fn=store_feedback,
+            inputs=[rating, feedback_text, chatbot, target_language],
         )
-
 if __name__ == "__main__":
     demo.launch()