Update app.py

app.py

@@ -3,8 +3,13 @@ import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from peft import PeftModel
 
+
 base_model_name = "facebook/nllb-200-distilled-600M"
-adapter_model_name = "entropy25/mt_en_no_oil"
+
+
+adapter_en_to_no = "entropy25/mt_en_no_oil"
+adapter_no_to_en = "entropy25/mt_no_en_oil"
+
 
 tokenizer = AutoTokenizer.from_pretrained(base_model_name)
 base_model = AutoModelForSeq2SeqLM.from_pretrained(
@@ -14,7 +19,8 @@ base_model = AutoModelForSeq2SeqLM.from_pretrained(
     device_map="auto"
 )
 
-model = PeftModel.from_pretrained(base_model, adapter_model_name)
+model_en_to_no = PeftModel.from_pretrained(base_model, adapter_en_to_no)
+model_no_to_en = PeftModel.from_pretrained(base_model, adapter_no_to_en)
 
 def translate(text, source_lang, target_lang):
     if not text.strip():
@@ -23,11 +29,18 @@ def translate(text, source_lang, target_lang):
     if source_lang == target_lang:
         return text
     
-    lang_map = {
-        "English": "eng_Latn",
-        "Norwegian": "nob_Latn"
-    }
+    if source_lang == "English" and target_lang == "Norwegian":
+        model = model_en_to_no
+        src_code = "eng_Latn"
+        tgt_code = "nob_Latn"
+    elif source_lang == "Norwegian" and target_lang == "English":
+        model = model_no_to_en
+        src_code = "nob_Latn"
+        tgt_code = "eng_Latn"
+    else:
+        return "Unsupported language pair"
     
+ 
     sentences = text.split('\n')
     translated_sentences = []
     
@@ -48,7 +61,7 @@ def translate(text, source_lang, target_lang):
         
         outputs = model.generate(
             **inputs,
-            forced_bos_token_id=tokenizer.convert_tokens_to_ids(lang_map[target_lang]),
+            forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_code),
             max_length=512,
             num_beams=5
         )
@@ -212,7 +225,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
     
     gr.HTML(
         "<div class='footer-info'>"
-        "Oil & Gas Translation • English ↔ Norwegian"
+        "Oil & Gas Translation • English ↔ Norwegian • Bidirectional Model"
         "</div>"
     )