Update app.py

app.py

@@ -6,15 +6,39 @@ from threading import Thread
 import time
 import pytz
 from datetime import datetime
+import gradio as gr
+import torch
+import time
+import pytz
+from datetime import datetime
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
 print("Loading model and tokenizer...")
 model_name = "large-traversaal/Phi-4-Hindi"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
 print("Model and tokenizer loaded successfully!")
-def generate_response(message, temperature, max_new_tokens, top_p):
-    print(f"Input: {message}")
+option_mapping = {
+    "translation": "### TRANSLATION ###",
+    "mcq": "### MCQ ###",
+    "nli": "### NLI ###",
+    "summarization": "### SUMMARIZATION ###",
+    "long response": "### LONG RESPONSE ###",
+    "short response": "### SHORT RESPONSE ###",
+    "direct response": "### DIRECT RESPONSE ###",
+    "paraphrase": "### PARAPHRASE ###",
+    "code": "### CODE ###"
+}
+def generate_response(message, temperature, max_new_tokens, top_p, task):
+    append_text = option_mapping.get(task, "")
+    prompt = f"INPUT : {message} {append_text} RESPONSE : "
+    print(f"Prompt: {prompt}")
     start_time = time.time()
-    inputs = tokenizer(message, return_tensors="pt").to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
     gen_kwargs = {
         "input_ids": inputs["input_ids"],
@@ -29,16 +53,12 @@ def generate_response(message, temperature, max_new_tokens, top_p):
     result = []
     for text in streamer:
         result.append(text)
-        current_output = "".join(result)
-        if current_output.startswith(message):
-            yield current_output[len(message):]
-        else:
-            yield current_output
+        yield "".join(result)
     end_time = time.time()
     time_taken = end_time - start_time
     output_text = "".join(result)
-    if output_text.startswith(message):
-        output_text = output_text[len(message):]
+    if "RESPONSE : " in output_text:
+        output_text = output_text.split("RESPONSE : ", 1)[1].strip()
     print(f"Output: {output_text}")
     print(f"Time taken: {time_taken:.2f} seconds")
     pst_timezone = pytz.timezone('America/Los_Angeles')
@@ -53,6 +73,11 @@ with gr.Blocks() as demo:
                 placeholder="Enter your text here...",
                 lines=5
             )
+            task_dropdown = gr.Dropdown(
+                choices=["translation", "mcq", "nli", "summarization", "long response", "short response", "direct response", "paraphrase", "code"],
+                value="long response",
+                label="Task"
+            )
             with gr.Row():
                 with gr.Column():
                     temperature = gr.Slider(
@@ -88,11 +113,11 @@ with gr.Blocks() as demo:
             )
     send_btn.click(
         fn=generate_response,
-        inputs=[input_text, temperature, max_new_tokens, top_p],
+        inputs=[input_text, temperature, max_new_tokens, top_p, task_dropdown],
         outputs=output_text
     )
     clear_btn.click(
-        fn=lambda: ("", "", "", ""),
+        fn=lambda: ("", ""),
         inputs=None,
         outputs=[input_text, output_text]
     )