Update app.py

app.py

@@ -3,84 +3,72 @@ import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from peft import PeftModel
 
-# 基础模型
 base_model_name = "facebook/nllb-200-distilled-600M"
-
-# 两个适配器模型
 adapter_en_to_no = "entropy25/mt_en_no_oil"
 adapter_no_to_en = "entropy25/mt_no_en_oil"
 
-# 加载 tokenizer
 tokenizer = AutoTokenizer.from_pretrained(base_model_name)
 
-# 方法1：为每个方向加载独立的基础模型
-print("Loading English to Norwegian model...")
-base_model_en_to_no = AutoModelForSeq2SeqLM.from_pretrained(
+print("Loading shared base model...")
+base_model = AutoModelForSeq2SeqLM.from_pretrained(
     base_model_name,
     torch_dtype=torch.float16,
     low_cpu_mem_usage=True,
     device_map="auto"
 )
-model_en_to_no = PeftModel.from_pretrained(base_model_en_to_no, adapter_en_to_no)
 
-print("Loading Norwegian to English model...")
-base_model_no_to_en = AutoModelForSeq2SeqLM.from_pretrained(
-    base_model_name,
-    torch_dtype=torch.float16,
-    low_cpu_mem_usage=True,
-    device_map="auto"
-)
-model_no_to_en = PeftModel.from_pretrained(base_model_no_to_en, adapter_no_to_en)
+print("Loading adapters...")
+model = PeftModel.from_pretrained(base_model, adapter_en_to_no, adapter_name="en_to_no")
+model.load_adapter(adapter_no_to_en, adapter_name="no_to_en")
 
 def translate(text, source_lang, target_lang):
-    if not text.strip():
-        return ""
-    
-    if source_lang == target_lang:
+    if not text.strip() or source_lang == target_lang:
         return text
     
-    # 选择正确的模型和语言代码
     if source_lang == "English" and target_lang == "Norwegian":
-        model = model_en_to_no
-        src_code = "eng_Latn"
-        tgt_code = "nob_Latn"
+        model.set_adapter("en_to_no")
+        src_code, tgt_code = "eng_Latn", "nob_Latn"
     elif source_lang == "Norwegian" and target_lang == "English":
-        model = model_no_to_en
-        src_code = "nob_Latn"
-        tgt_code = "eng_Latn"
+        model.set_adapter("no_to_en")
+        src_code, tgt_code = "nob_Latn", "eng_Latn"
     else:
         return "Unsupported language pair"
     
-    # 按行翻译
-    sentences = text.split('\n')
-    translated_sentences = []
+    lines = text.split('\n')
+    non_empty_lines = [line for line in lines if line.strip()]
     
-    for sentence in sentences:
-        if not sentence.strip():
-            translated_sentences.append("")
-            continue
-        
-        inputs = tokenizer(
-            sentence,
-            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(tgt_code),
-            max_length=512,
-            num_beams=5
-        )
-        
-        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        translated_sentences.append(result)
+    if not non_empty_lines:
+        return text
+    
+    inputs = tokenizer(
+        non_empty_lines,
+        return_tensors="pt",
+        padding=True,
+        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(tgt_code),
+        max_length=512,
+        num_beams=3
+    )
+    
+    results = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    
+    result_iter = iter(results)
+    final_lines = []
+    for line in lines:
+        if line.strip():
+            final_lines.append(next(result_iter))
+        else:
+            final_lines.append("")
     
-    return '\n'.join(translated_sentences)
+    return '\n'.join(final_lines)
 
 def swap_languages(src, tgt, input_txt, output_txt):
     return tgt, src, output_txt, input_txt
@@ -88,7 +76,6 @@ def swap_languages(src, tgt, input_txt, output_txt):
 def load_file(file):
     if file is None:
         return ""
-    
     try:
         with open(file.name, 'r', encoding='utf-8') as f:
             return f.read()
@@ -99,54 +86,43 @@ def load_file(file):
         except Exception as e:
             return f"Error reading file: {str(e)}"
 
-# 英语示例
 EXAMPLES_EN = {
     "drilling_short": "Mud weight adjusted to 1.82 specific gravity at 3,247 meters depth.",
     "drilling_long": "The drilling operation at well site A-15 encountered unexpected high-pressure zones at 3,247 meters depth, requiring immediate adjustment of mud weight from 1.65 to 1.82 specific gravity to maintain wellbore stability and prevent potential kicks.",
-    
     "reservoir_short": "Permeability is 250 millidarcy with 22 percent porosity.",
     "reservoir_long": "The reservoir shows excellent permeability of 250 millidarcy and porosity of 22 percent based on core analysis, indicating significant hydrocarbon potential with estimated oil saturation of 65 percent.",
-    
     "subsea_short": "Christmas tree rated for 10,000 psi working pressure.",
     "subsea_long": "The subsea production system consists of a vertical Christmas tree rated for 10,000 psi working pressure and 150 degrees Celsius temperature, equipped with redundant safety features including automatic shutdown valves and real-time pressure monitoring systems.",
-    
     "seismic_short": "Structural trap area estimated at 12 square kilometers.",
     "seismic_long": "Seismic data confirms the presence of a structural trap with an estimated area of 12 square kilometers, and productivity tests show stabilized oil production of 3,400 barrels per day at optimization pressure of 2,100 psi.",
-    
     "safety_short": "H2S training required before site access.",
     "safety_long": "Emergency response procedures require all personnel to complete H2S safety training before site access, with breathing apparatus and wind indicators positioned at designated muster points, and immediate evacuation protocols activated when gas detection exceeds 10 ppm concentration levels."
 }
 
-# 挪威语示例
 EXAMPLES_NO = {
     "drilling_short": "Slamvekt justert til 1,82 spesifikk tyngde ved 3 247 meters dybde.",
     "drilling_long": "Boreoperasjonen ved brønnsted A-15 støtte på uventede høytrykksoner ved 3 247 meters dybde, noe som krevde umiddelbar justering av slamvekt fra 1,65 til 1,82 spesifikk tyngde for å opprettholde brønnborestabilitet og forhindre potensielle kicks.",
-    
     "reservoir_short": "Permeabilitet er 250 millidarcy med 22 prosent porøsitet.",
     "reservoir_long": "Reservoaret viser utmerket permeabilitet på 250 millidarcy og porøsitet på 22 prosent basert på kjerneanalyse, noe som indikerer betydelig hydrokarbonpotensial med estimert oljemetning på 65 prosent.",
-    
     "subsea_short": "Juletre dimensjonert for 10 000 psi arbeidstrykk.",
     "subsea_long": "Subsea produksjonssystemet består av et vertikalt juletre dimensjonert for 10 000 psi arbeidstrykk og 150 grader Celsius temperatur, utstyrt med redundante sikkerhetsfunksjoner inkludert automatiske nedstengningsventiler og sanntids trykkmonitorering.",
-    
     "seismic_short": "Strukturell felle estimert til 12 kvadratkilometer.",
     "seismic_long": "Seismiske data bekrefter tilstedeværelsen av en strukturell felle med estimert areal på 12 kvadratkilometer, og produktivitetstester viser stabilisert oljeproduksjon på 3 400 fat per dag ved optimaliseringstrykk på 2 100 psi.",
-    
     "safety_short": "H2S-opplæring påkrevd før tilgang til området.",
     "safety_long": "Nødprosedyrer krever at alt personell fullfører H2S-sikkerhetsopplæring før områdetilgang, med åndedrettsutstyr og vindindikatorer plassert ved utpekte samlingspunkter, og umiddelbare evakueringsprotokoller aktiveres når gassdeteksjon overskrider 10 ppm konsentrasjonsnivå."
 }
 
 def get_examples(source_lang):
-    """根据源语言返回对应的示例"""
-    if source_lang == "English":
-        return EXAMPLES_EN
-    else:
-        return EXAMPLES_NO
+    return EXAMPLES_EN if source_lang == "English" else EXAMPLES_NO
 
 def update_example_buttons(source_lang):
-    """当语言改变时更新示例"""
     examples = get_examples(source_lang)
     return examples["drilling_short"]
 
+def get_example(key, source_lang):
+    examples = get_examples(source_lang)
+    return examples[key]
+
 custom_css = """
 .gradio-container {
     max-width: 1100px !important;
@@ -214,7 +190,6 @@ custom_css = """
 """
 
 with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
-    
     gr.HTML("<div style='height: 20px'></div>")
     
     with gr.Row():
@@ -229,7 +204,6 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
                         elem_classes="lang-selector",
                         scale=1
                     )
-                
                 input_text = gr.Textbox(
                     placeholder="Type to translate",
                     show_label=False,
@@ -254,7 +228,6 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
                         elem_classes="lang-selector",
                         scale=1
                     )
-                
                 output_text = gr.Textbox(
                     placeholder="Translation",
                     show_label=False,
@@ -265,11 +238,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
                     interactive=False
                 )
     
-    gr.HTML(
-        "<div class='footer-info'>"
-        "Oil & Gas Translation • English ↔ Norwegian • Bidirectional Model"
-        "</div>"
-    )
+    gr.HTML("<div class='footer-info'>Oil & Gas Translation • English ↔ Norwegian • Bidirectional Model</div>")
     
     with gr.Accordion("Example Sentences", open=True):
         with gr.Row():
@@ -296,11 +265,6 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
             btn9 = gr.Button("Safety (Short)", size="sm")
             btn10 = gr.Button("Safety (Long)", size="sm")
         
-        # 为每个按钮设置点击事件,根据当前源语言选择示例
-        def get_example(key, source_lang):
-            examples = get_examples(source_lang)
-            return examples[key]
-        
         btn1.click(lambda sl: get_example("drilling_short", sl), inputs=[source_lang], outputs=example_text)
         btn2.click(lambda sl: get_example("drilling_long", sl), inputs=[source_lang], outputs=example_text)
         btn3.click(lambda sl: get_example("reservoir_short", sl), inputs=[source_lang], outputs=example_text)
@@ -312,11 +276,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
         btn9.click(lambda sl: get_example("safety_short", sl), inputs=[source_lang], outputs=example_text)
         btn10.click(lambda sl: get_example("safety_long", sl), inputs=[source_lang], outputs=example_text)
         
-        use_example_btn.click(
-            fn=lambda x: x,
-            inputs=example_text,
-            outputs=input_text
-        )
+        use_example_btn.click(fn=lambda x: x, inputs=example_text, outputs=input_text)
     
     with gr.Accordion("Upload Text File", open=False):
         file_input = gr.File(
@@ -325,41 +285,11 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
             type="filepath"
         )
     
-    # 当源语言改变时,更新示例显示
-    source_lang.change(
-        fn=update_example_buttons,
-        inputs=[source_lang],
-        outputs=[example_text]
-    )
-    
-    input_text.change(
-        fn=translate,
-        inputs=[input_text, source_lang, target_lang],
-        outputs=output_text
-    )
-    
-    source_lang.change(
-        fn=translate,
-        inputs=[input_text, source_lang, target_lang],
-        outputs=output_text
-    )
-    
-    target_lang.change(
-        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, output_text],
-        outputs=[source_lang, target_lang, input_text, output_text]
-    )
-    
-    file_input.change(
-        fn=load_file,
-        inputs=file_input,
-        outputs=input_text
-    )
+    source_lang.change(fn=update_example_buttons, inputs=[source_lang], outputs=[example_text])
+    input_text.change(fn=translate, inputs=[input_text, source_lang, target_lang], outputs=output_text)
+    source_lang.change(fn=translate, inputs=[input_text, source_lang, target_lang], outputs=output_text)
+    target_lang.change(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, output_text], outputs=[source_lang, target_lang, input_text, output_text])
+    file_input.change(fn=load_file, inputs=file_input, outputs=input_text)
 
 demo.launch()