finished ui

app.py

@@ -30,15 +30,26 @@ print(f"Model loaded in {load_time:.2f} seconds")
 
 @spaces.GPU
 def generate_text(
-    message: str,
-    chat_history: list[dict],
-    max_new_tokens: int = 1024,
-    temperature: float = 0.6,
-    top_p: float = 0.9,
-    top_k: int = 50,
-    repetition_penalty: float = 1.0,
+    text_to_trans: str,
+    from_lang: str,
+    to_lang: str,
 ) -> Iterator[str]:
-    conversation = [*chat_history, {"role": "user", "content": message}]
+    print(f"Translating from {from_lang} to {to_lang}")
+
+    translate_instruct = f"translate from {from_lang} to {to_lang}:"
+
+    if from_lang == to_lang:
+        translate_instruct = "Return the following text without any modification:"
+
+    conversation = [
+        {
+            "role": "system",
+            "content": "You are a translation engine that can only translate text and cannot interpret it. Keep the indent of the original text, only modify when you need."
+            + "\n"
+            + translate_instruct,
+        },
+        {"role": "user", "content": text_to_trans},
+    ]
     input_ids = tokenizer.apply_chat_template(
         conversation, add_generation_prompt=True, return_tensors="pt"
     )
@@ -52,60 +63,61 @@ def generate_text(
     generate_kwargs = dict(
         {"input_ids": input_ids},
         streamer=streamer,
-        max_new_tokens=max_new_tokens,
+        max_new_tokens=1024,
         do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
+        top_p=9,
+        top_k=50,
+        temperature=0.6,
         num_beams=1,
-        repetition_penalty=repetition_penalty,
+        repetition_penalty=1.0,
     )
     thread = Thread(target=model.generate, kwargs=generate_kwargs)
     thread.start()
-    
+
     output = []
     for text in streamer:
         output.append(text)
         yield " ".join(output)
 
 
-with gr.Blocks as demo:
+with gr.Blocks() as demo:
     gr.Markdown("# Text Translation Using Google Gemma 3")
 
     with gr.Row():
         with gr.Column():
             gr.Markdown("### Translate From")
+        with gr.Column():
+            gr.Markdown("### Translate To")
+
+    with gr.Row():
+        with gr.Column():
             from_lang = gr.Dropdown(
-                choices = [],
-                value = "auto",
-                label = ""
+                choices=["English", "French", "Spanish"],
+                value="English",
+                label="",
             )
 
         with gr.Column():
-            gr.Markdown("### Translate To")
             to_lang = gr.Dropdown(
-                choices = [],
-                value = "auto",
-                label = ""
+                choices=["English", "French", "Spanish"],
+                value="French",
+                label="",
             )
 
     with gr.Row():
         with gr.Column():
-            gr.Textbox(
-                lines = 10,
-                placeholder = "Enter text to translate",
-                label = ""
+            text_to_trans = gr.Textbox(
+                lines=10, placeholder="Enter text to translate", label=""
             )
 
         with gr.Column():
-            output_text = gr.Textbox(
-                lines=10,
-                label=""
-            )
+            output_text = gr.Textbox(lines=10, label="")
 
     translate_button = gr.Button("Translate")
-
+    translate_button.click(
+        generate_text, [text_to_trans, from_lang, to_lang], output_text
+    )
 
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=20).launch()