Update app.py

app.py

@@ -3,62 +3,142 @@ import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from peft import PeftModel
 
-# Base model
-base_model_name = "facebook/nllb-200-distilled-600M"
-# LoRA adapter
-adapter_model_name = "entropy25/mt_en_no_oil"
+BASE_MODEL = "facebook/nllb-200-distilled-600M"
+ADAPTER_NO_TO_EN = "entropy25/mt_en_no_oil"
+#ADAPTER_EN_TO_NO = "entropy25/no_en"
 
-# Load with optimizations to reduce storage
-tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 base_model = AutoModelForSeq2SeqLM.from_pretrained(
-    base_model_name,
-    torch_dtype=torch.float16,  # Use half precision
+    BASE_MODEL,
+    torch_dtype=torch.float16,
     low_cpu_mem_usage=True,
     device_map="auto"
 )
 
-# Apply the LoRA adapter
-model = PeftModel.from_pretrained(base_model, adapter_model_name)
+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"
+}
 
 def translate(text, source_lang, target_lang):
     if not text.strip():
-        return "Please enter text to translate."
+        return "Please enter text to translate"
+    
+    if source_lang == target_lang:
+        return "Source and target languages must be different"
+    
+    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
+    
+    except Exception as e:
+        return f"Translation error: {str(e)}"
+
+def swap_languages(source, target, text, translation):
+    return target, source, translation, text
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# Oil & Gas Professional Translation")
+    gr.Markdown("English ↔ Norwegian translation specialized for petroleum industry")
+    
+    with gr.Row():
+        source_lang = gr.Dropdown(
+            choices=["English", "Norwegian"],
+            label="Source Language",
+            value="English"
+        )
+        
+        swap_btn = gr.Button("⇄", scale=0, size="sm")
+        
+        target_lang = gr.Dropdown(
+            choices=["English", "Norwegian"],
+            label="Target Language",
+            value="Norwegian"
+        )
     
-    lang_map = {
-        "English": "eng_Latn",
-        "Norwegian": "nob_Latn"
-    }
+    with gr.Row():
+        with gr.Column():
+            input_text = gr.Textbox(
+                label="Input Text",
+                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
+            )
+            with gr.Row():
+                copy_btn = gr.Button("📋 Copy", scale=1)
+                clear_btn = gr.Button("🗑️ Clear", scale=1)
     
-    inputs = tokenizer(
-        text,
-        return_tensors="pt",
-        truncation=True,
-        max_length=512
+    translate_btn = gr.Button("Translate", variant="primary", size="lg")
+    
+    gr.Examples(
+        examples=[
+            ["The drilling operation encountered high pressure", "English", "Norwegian"],
+            ["Reservoaret viser god permeabilitet", "Norwegian", "English"]
+        ],
+        inputs=[input_text, source_lang, target_lang]
+    )
+    
+    input_text.change(
+        fn=lambda x: str(len(x)),
+        inputs=input_text,
+        outputs=input_chars
     )
     
-    # Move inputs to the same device as model
-    if hasattr(model, 'device'):
-        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    translate_btn.click(
+        fn=translate,
+        inputs=[input_text, source_lang, target_lang],
+        outputs=output_text
+    )
     
-    outputs = model.generate(
-        **inputs,
-        forced_bos_token_id=tokenizer.convert_tokens_to_ids(lang_map[target_lang]),
-        max_length=512,
-        num_beams=5
+    swap_btn.click(
+        fn=swap_languages,
+        inputs=[source_lang, target_lang, input_text, output_text],
+        outputs=[source_lang, target_lang, input_text, output_text]
     )
     
-    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return result
+    copy_btn.click(
+        fn=lambda x: x,
+        inputs=output_text,
+        outputs=input_text
+    )
+    
+    clear_btn.click(
+        fn=lambda: ("", ""),
+        outputs=[input_text, output_text]
+    )
 
-# Simple Gradio UI
-gr.Interface(
-    fn=lambda text, src, tgt: translate(text, src, tgt),
-    inputs=[
-        gr.Textbox(label="Input text", lines=6),
-        gr.Dropdown(choices=["English", "Norwegian"], label="Source language", value="English"),
-        gr.Dropdown(choices=["English", "Norwegian"], label="Target language", value="Norwegian")
-    ],
-    outputs=gr.Textbox(label="Translation", lines=6),
-    title="LoRA-Enhanced English↔Norwegian Translator",
-    description="Fine-tuned NLLB-200 model with LoRA adapter: entropy25/mt_en_no_oil"
-).launch()
+demo.launch()