Update app.py

app.py

@@ -1,143 +1,91 @@
-import tkinter as tk
-from tkinter import ttk
+import gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
+# -----------------------
+# Translation core
+# -----------------------
 class TranslationPipeline:
-    def __init__(self, model_name):
+    def __init__(self, model_name: str, device: str = "cpu"):
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        self.device = device
+        self.model.to(self.device)
 
-    def __call__(self, text):
-        inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-        outputs = self.model.generate(**inputs)
+    @torch.inference_mode()
+    def __call__(self, text: str) -> str:
+        inputs = self.tokenizer(
+            text, return_tensors="pt", padding=True, truncation=True
+        ).to(self.device)
+
+        outputs = self.model.generate(**inputs, max_new_tokens=256)
         translated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return [{'translation_text': translated_text}]
-
-class TranslateApp:
-    def __init__(self, root):
-        self.root = root
-        self.root.title("Short Translation")
-        self.root.geometry("700x300")
-        self.root.configure(bg="#f0f0f0")
-
-        # Define Colors
-        self.maroon = "#800000"
-        self.white  = "#ffffff"
-
-        # Initialize Translation Models (pipeline parameters)
-        self.models = {
-            "Spanish": ["translation_en_to_es", 
-                        "Helsinki-NLP/opus-mt-en-es"],
-            "German":  ["translation_en_to_de",
-                        "Helsinki-NLP/opus-mt-en-de"],
-            "Japanese": ["translation_en_to_ja",
-                         "staka/fugumt-en-ja"],
-            "Ukrainian": ["translation_en_to_uk",
-                         "Helsinki-NLP/opus-mt-en-uk"],
-            "Russian":   ["translation_en_to_ru",
-                         "Helsinki-NLP/opus-mt-en-ru"],
-        }
-        
-        # Cache for loaded pipeline objects
-        # This is done to speed up translations after the first.
-        self.cached_pipelines = {}
-
-        # --- Top Section ---
-        top_frame = tk.Frame(self.root, bg="#f0f0f0")
-        top_frame.pack(side=tk.TOP, fill=tk.X, padx=10, pady=10)
-
-        # Left Column: Input and Buttons
-        left_column = tk.Frame(top_frame, bg="#f0f0f0", borderwidth=2, relief="groove")
-        left_column.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=(0, 5))
-
-        tk.Label(left_column, text="English Sentence", fg=self.maroon, 
-            bg="#f0f0f0", font=("Arial", 12, "bold")).pack(anchor="w", padx=5)
-        
-        self.input_entry = tk.Entry(left_column, font=("Arial", 12))
-        self.input_entry.pack(fill=tk.X, padx=10, pady=5)
-
-        button_frame = tk.Frame(left_column, bg="#f0f0f0")
-        button_frame.pack(pady=5)
-
-        self.translate_btn = tk.Button(button_frame, text="Translate", 
-                                       fg=self.maroon, bg=self.white, 
-                                       command=self.translate_text, width=15)
-        self.translate_btn.pack(side=tk.LEFT, padx=5)
-
-        self.clear_btn = tk.Button(button_frame, text="Clear", fg=self.maroon, 
-                                   bg=self.white, 
-                                   command=self.clear_fields, width=15)
-        self.clear_btn.pack(side=tk.LEFT, padx=5)
-
-        # Right Column: Language Selection
-        right_column = tk.Frame(top_frame, bg="#f0f0f0", borderwidth=2, 
-                                relief="groove")
-        right_column.pack(side=tk.RIGHT, fill=tk.Y, padx=(5, 0))
-
-        tk.Label(right_column, text="Translation Language", fg=self.maroon, 
-            bg="#f0f0f0", font=("Arial", 12, "bold")).pack(anchor="w", padx=5)
-
-        self.lang_var = tk.StringVar(value="Spanish")
-        languages = [("Spanish", "es"), ("German", "de"), ("Japanese", "ja"), 
-                     ("Ukrainian", "uk"), ("Russian", "ru")]
-        
-        for lang_text, lang_code in languages:
-            tk.Radiobutton(right_column, text=lang_text,
-                        font=("Arial", 12, "bold"),
-                        variable=self.lang_var, value=lang_text,
-                        fg=self.maroon, bg="#f0f0f0", highlightthickness=0, 
-                        activeforeground=self.maroon).pack(anchor="w", padx=30)
-
-        # --- Bottom Section: Output Area ---
-        self.output_frame = tk.Frame(self.root, bg=self.maroon, height=120)
-        self.output_frame.pack(side=tk.TOP, fill=tk.X, padx=10, pady=(0, 10))
-        self.output_frame.pack_propagate(False) # Maintain fixed height
-
-        self.output_label = tk.Label(self.output_frame, text="", fg=self.white, 
-                                     bg=self.maroon, 
-                                     font=("Times New Roman", 16, "bold"), 
-                                     wraplength=650, justify="center")
-        self.output_label.pack(expand=True, fill=tk.BOTH)
-
-    def translate_text(self):
-        input_text  = self.input_entry.get().strip()
-        target_lang = self.lang_var.get()
-
-        if not input_text:
-            self.output_label.config(text="Please enter text to translate.")
-            return
-
-        # Check if we need to load the model
-        if target_lang not in self.cached_pipelines:
-            self.output_label.config(text=f"Loading model for {target_lang}...")
-            self.root.update_idletasks()
-            
-            try:
-                # We ignore the task name (index 0) since we are using the custom pipeline
-                _, model_name = self.models[target_lang]
-                self.cached_pipelines[target_lang] = TranslationPipeline(model_name)
-            except Exception as e:
-                self.output_label.config(text=f"Error loading model: {str(e)}")
-                return
-
-        self.output_label.config(text=f"Translating to {target_lang}...")
-        self.root.update_idletasks()
-
-        try:
-            translator = self.cached_pipelines[target_lang]
-            result     = translator(input_text)
-            translated_text = result[0]['translation_text']
-
-            print(f"DEBUG: {target_lang} output -> {translated_text}")
-            self.output_label.config(text=translated_text)
-        except Exception as e:
-            self.output_label.config(text=f"Error: {str(e)}")
-
-    def clear_fields(self):
-        self.input_entry.delete(0, tk.END)
-        self.output_label.config(text="")
+        return translated_text
+
+
+# -----------------------
+# Models + cache
+# -----------------------
+MODELS = {
+    "Spanish": "Helsinki-NLP/opus-mt-en-es",
+    "German": "Helsinki-NLP/opus-mt-en-de",
+    "Japanese": "staka/fugumt-en-ja",
+    "Ukrainian": "Helsinki-NLP/opus-mt-en-uk",
+    "Russian": "Helsinki-NLP/opus-mt-en-ru",
+}
+
+# Cache loaded pipelines so we don’t re-download/reload every time
+PIPELINE_CACHE = {}
+
+# Use GPU if available (some Spaces have it; many are CPU)
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def translate(text: str, target_lang: str) -> str:
+    text = (text or "").strip()
+    if not text:
+        return "Please enter text to translate."
+
+    if target_lang not in MODELS:
+        return "Unsupported language selection."
+
+    if target_lang not in PIPELINE_CACHE:
+        model_name = MODELS[target_lang]
+        # Loading can take time on first request
+        PIPELINE_CACHE[target_lang] = TranslationPipeline(model_name, device=DEVICE)
+
+    translator = PIPELINE_CACHE[target_lang]
+    return translator(text)
+
+
+# -----------------------
+# Gradio UI
+# -----------------------
+with gr.Blocks(title="Short Translation") as demo:
+    gr.Markdown("## Short Translation\nEnter an English sentence and choose a target language.")
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            input_text = gr.Textbox(label="English Sentence", lines=3, placeholder="Type here...")
+            translate_btn = gr.Button("Translate")
+            clear_btn = gr.Button("Clear")
+
+        with gr.Column(scale=1):
+            target_lang = gr.Radio(
+                choices=list(MODELS.keys()),
+                value="Spanish",
+                label="Translation Language",
+            )
+
+    output_text = gr.Textbox(label="Translation", lines=4)
+
+    translate_btn.click(fn=translate, inputs=[input_text, target_lang], outputs=output_text)
+
+    def clear():
+        return "", ""
+
+    clear_btn.click(fn=clear, inputs=None, outputs=[input_text, output_text])
 
 if __name__ == "__main__":
-    root = tk.Tk()
-    app = TranslateApp(root)
-    root.mainloop()
+    demo.launch()
+