Update app.py

app.py

@@ -1,46 +1,22 @@
-import gradio as gr
-import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from peft import PeftModel
-import os
-import shutil
-import gc
-
-def cleanup_cache():
-    cache_dirs = ["/tmp/cache", "/tmp/hf", "~/.cache/huggingface"]
-    for cache_dir in cache_dirs:
-        expanded_dir = os.path.expanduser(cache_dir)
-        if os.path.exists(expanded_dir):
-            try:
-                shutil.rmtree(expanded_dir)
-                print(f"Cleaned: {expanded_dir}")
-            except:
-                pass
-    gc.collect()
-    torch.cuda.empty_cache() if torch.cuda.is_available() else None
-
-cleanup_cache()
 
 BASE_MODEL = "facebook/nllb-200-distilled-600M"
-ADAPTER_MODEL = "entropy25/mt_en_no_oil"
+ADAPTER_NO_TO_EN = "entropy25/mt_en_no_oil"
+#ADAPTER_EN_TO_NO = "entropy25/no_en"
+
 
-print("Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 
-print("Loading model with aggressive optimization...")
 base_model = AutoModelForSeq2SeqLM.from_pretrained(
     BASE_MODEL,
     torch_dtype=torch.float16,
     low_cpu_mem_usage=True,
-    device_map="auto",
-    use_safetensors=True
+    device_map="auto"
 )
 
-print("Loading LoRA adapter...")
-model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL)
-
-print("Model ready. Cleaning up...")
-cleanup_cache()
+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",
@@ -55,6 +31,8 @@ def translate(text, 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",
@@ -83,7 +61,7 @@ def swap_languages(source, target, text, translation):
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# Oil & Gas Professional Translation")
-    gr.Markdown("English ↔ Norwegian specialized for petroleum industry")
+    gr.Markdown("English ↔ Norwegian translation specialized for petroleum industry")
     
     with gr.Row():
         source_lang = gr.Dropdown(
@@ -150,6 +128,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         fn=swap_languages,
         inputs=[source_lang, target_lang, input_text, output_text],
         outputs=[source_lang, target_lang, input_text, output_text]
+
     )
     
     copy_btn.click(