Update app.py

app.py

@@ -1,63 +1,60 @@
+import gradio as gr
+import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from peft import PeftModel
 
-BASE_MODEL = "facebook/nllb-200-distilled-600M"
-ADAPTER_NO_TO_EN = "entropy25/mt_en_no_oil"
-#ADAPTER_EN_TO_NO = "entropy25/no_en"
-
-
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+base_model_name = "facebook/nllb-200-distilled-600M"
+adapter_model_name = "entropy25/mt_en_no_oil"
 
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
 base_model = AutoModelForSeq2SeqLM.from_pretrained(
-    BASE_MODEL,
+    base_model_name,
     torch_dtype=torch.float16,
     low_cpu_mem_usage=True,
     device_map="auto"
 )
 
-model_no_to_en = PeftModel.from_pretrained(base_model, ADAPTER_NO_TO_EN)
-model_en_to_no = PeftModel.from_pretrained(base_model, ADAPTER_EN_TO_NO)
-
-LANG_CODES = {
-    "English": "eng_Latn",
-    "Norwegian": "nob_Latn"
-}
+model = PeftModel.from_pretrained(base_model, adapter_model_name)
 
 def translate(text, source_lang, target_lang):
     if not text.strip():
-        return "Please enter text to translate"
+        return "Please enter text to translate", "0 characters"
     
     if source_lang == target_lang:
-        return "Source and target languages must be different"
+        return "Source and target languages cannot be the same", f"{len(text)} characters"
     
-    try:
-        model = model_no_to_en if source_lang == "Norwegian" else model_en_to_no
-        
-        inputs = tokenizer(
-            text,
-            return_tensors="pt",
-            truncation=True,
-            max_length=512
-        )
-        
-        if hasattr(model, 'device'):
-            inputs = {k: v.to(model.device) for k, v in inputs.items()}
-        
-        outputs = model.generate(
-            **inputs,
-            forced_bos_token_id=tokenizer.convert_tokens_to_ids(LANG_CODES[target_lang]),
-            max_length=512,
-            num_beams=5
-        )
-        
-        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return result
+    lang_map = {
+        "English": "eng_Latn",
+        "Norwegian": "nob_Latn"
+    }
+    
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        max_length=512
+    )
+    
+    if hasattr(model, 'device'):
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    
+    outputs = model.generate(
+        **inputs,
+        forced_bos_token_id=tokenizer.convert_tokens_to_ids(lang_map[target_lang]),
+        max_length=512,
+        num_beams=5
+    )
     
-    except Exception as e:
-        return f"Translation error: {str(e)}"
+    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    char_info = f"{len(text)} → {len(result)} characters"
+    
+    return result, char_info
+
+def swap_languages(src, tgt):
+    return tgt, src
 
-def swap_languages(source, target, text, translation):
-    return target, source, translation, text
+def clear_all():
+    return "", "", "0 characters"
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# Oil & Gas Professional Translation")
@@ -69,9 +66,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             label="Source Language",
             value="English"
         )
-        
-        swap_btn = gr.Button("⇄", scale=0, size="sm")
-        
+        swap_btn = gr.Button("⇄ Swap", size="sm")
         target_lang = gr.Dropdown(
             choices=["English", "Norwegian"],
             label="Target Language",
@@ -85,61 +80,44 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 placeholder="Enter text to translate",
                 lines=8
             )
-            input_chars = gr.Textbox(
-                label="Character Count",
-                value="0",
-                interactive=False,
-                max_lines=1
-            )
         
         with gr.Column():
             output_text = gr.Textbox(
                 label="Translation",
-                lines=8,
-                interactive=False
+                lines=8
             )
-            with gr.Row():
-                copy_btn = gr.Button("📋 Copy", scale=1)
-                clear_btn = gr.Button("🗑️ Clear", scale=1)
     
-    translate_btn = gr.Button("Translate", variant="primary", size="lg")
+    char_count = gr.Textbox(label="Statistics", value="0 characters", interactive=False)
+    
+    with gr.Row():
+        translate_btn = gr.Button("Translate", variant="primary", size="lg")
+        clear_btn = gr.Button("Clear All", size="lg")
     
     gr.Examples(
         examples=[
-            ["The drilling operation encountered high pressure", "English", "Norwegian"],
-            ["Reservoaret viser god permeabilitet", "Norwegian", "English"]
+            ["The drilling operation encountered high pressure at 3000 meters depth", "English", "Norwegian"],
+            ["Reservoaret viser god permeabilitet og porøsitet", "Norwegian", "English"],
+            ["Wellhead pressure monitoring indicates stable production", "English", "Norwegian"]
         ],
-        inputs=[input_text, source_lang, target_lang]
-    )
-    
-    input_text.change(
-        fn=lambda x: str(len(x)),
-        inputs=input_text,
-        outputs=input_chars
+        inputs=[input_text, source_lang, target_lang],
+        label="Example Translations"
     )
     
     translate_btn.click(
         fn=translate,
         inputs=[input_text, source_lang, target_lang],
-        outputs=output_text
+        outputs=[output_text, char_count]
     )
     
     swap_btn.click(
         fn=swap_languages,
-        inputs=[source_lang, target_lang, input_text, output_text],
-        outputs=[source_lang, target_lang, input_text, output_text]
-
-    )
-    
-    copy_btn.click(
-        fn=lambda x: x,
-        inputs=output_text,
-        outputs=input_text
+        inputs=[source_lang, target_lang],
+        outputs=[source_lang, target_lang]
     )
     
     clear_btn.click(
-        fn=lambda: ("", ""),
-        outputs=[input_text, output_text]
+        fn=clear_all,
+        outputs=[input_text, output_text, char_count]
     )
 
 demo.launch()