Update app.py

app.py

@@ -1,131 +1,59 @@
 import gradio as gr
 import torch
-from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from peft import PeftModel
 
-# Load Facebook's mBART-50 model
-model_name = "facebook/mbart-large-50-many-to-many-mmt"
-tokenizer = MBart50TokenizerFast.from_pretrained(model_name)
-model = MBartForConditionalGeneration.from_pretrained(model_name)
+# ① Base model（基础模型）
+base_model_name = "facebook/nllb-200-distilled-600M"
+
+# ② LoRA adapter（你的模型）
+adapter_model_name = "entropy25/mt_en_no_oil"
+
+# Load tokenizer and base model
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+base_model = AutoModelForSeq2SeqLM.from_pretrained(base_model_name)
+
+# Apply the LoRA adapter
+model = PeftModel.from_pretrained(base_model, adapter_model_name)
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
 def translate(text, source_lang, target_lang):
     if not text.strip():
-        return "Please enter text to translate"
+        return "Please enter text to translate."
     
-    # Language codes for mBART-50
     lang_map = {
-        "Norwegian": "nb_NO",
-        "English": "en_XX"
+        "English": "eng_Latn",
+        "Norwegian": "nob_Latn"
     }
     
-    try:
-        # Set source language
-        tokenizer.src_lang = lang_map[source_lang]
-        
-        # Tokenize
-        encoded = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
-        encoded = encoded.to(device)
-        
-        # Generate translation
-        generated = model.generate(
-            **encoded,
-            forced_bos_token_id=tokenizer.lang_code_to_id[lang_map[target_lang]],
-            max_length=512,
-            num_beams=5,
-            early_stopping=True
-        )
-        
-        # Decode result
-        result = tokenizer.batch_decode(generated, skip_special_tokens=True)[0]
-        return result
-        
-    except Exception as e:
-        return f"Translation error: {str(e)}"
-
-def swap_languages(source, target, text):
-    return target, source, ""
-
-# Create Gradio interface
-with gr.Blocks(
-    title="🇳🇴 ↔️ 🇬🇧 Norwegian-English Translator",
-    theme=gr.themes.Soft()
-) as demo:
-    
-    gr.Markdown("""
-    # 🌍 Advanced Norwegian-English Translator
-    ### Powered by Facebook's mBART-50 Model
-    """)
-    
-    with gr.Row():
-        source_lang = gr.Dropdown(
-            choices=["Norwegian", "English"],
-            value="Norwegian",
-            label="Source Language"
-        )
-        
-        swap_btn = gr.Button("🔄", size="sm")
-        
-        target_lang = gr.Dropdown(
-            choices=["Norwegian", "English"],
-            value="English",
-            label="Target Language"
-        )
-    
-    with gr.Row():
-        with gr.Column():
-            input_text = gr.Textbox(
-                lines=8,
-                placeholder="Enter text to translate...",
-                label="Input Text",
-                max_lines=15
-            )
-            
-        with gr.Column():
-            output_text = gr.Textbox(
-                lines=8,
-                label="Translation",
-                interactive=False,
-                max_lines=15
-            )
-    
-    translate_btn = gr.Button("🚀 Translate", variant="primary", size="lg")
-    
-    # Event handlers
-    translate_btn.click(
-        fn=translate,
-        inputs=[input_text, source_lang, target_lang],
-        outputs=output_text
-    )
-    
-    swap_btn.click(
-        fn=swap_languages,
-        inputs=[source_lang, target_lang, input_text],
-        outputs=[source_lang, target_lang, input_text]
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        max_length=512
+    ).to(device)
+    
+    outputs = model.generate(
+        **inputs,
+        forced_bos_token_id=tokenizer.convert_tokens_to_ids(lang_map[target_lang]),
+        max_length=512,
+        num_beams=5
     )
     
-    input_text.submit(
-        fn=translate,
-        inputs=[input_text, source_lang, target_lang],
-        outputs=output_text
-    )
-    
-    gr.Markdown("""
-    ### 📝 Examples
-    **Norwegian to English:**
-    - "Hei, hvordan har du det?" → "Hello, how are you?"
-    - "Jeg kommer fra Norge" → "I come from Norway"
-    
-    **English to Norwegian:**
-    - "Thank you very much" → "Tusen takk"
-    - "Good morning" → "God morgen"
-    """)
-    
-    gr.Markdown("""
-    ---
-    *Built with ❤️ using Facebook's mBART-50 and Gradio*
-    """)
+    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return result
 
-if __name__ == "__main__":
-    demo.launch(share=True)
+# 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(share=True)