Update app.py

app.py

@@ -3,11 +3,14 @@ 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)
 base_model = AutoModelForSeq2SeqLM.from_pretrained(
     base_model_name,
@@ -16,10 +19,10 @@ base_model = AutoModelForSeq2SeqLM.from_pretrained(
     device_map="auto"
 )
 
+# 加载两个方向的模型
 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():
         return ""
@@ -27,17 +30,19 @@ def translate(text, source_lang, target_lang):
     if 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"
     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 = []
     
@@ -46,7 +51,12 @@ def translate(text, source_lang, target_lang):
             translated_sentences.append("")
             continue
         
-        inputs = tokenizer(sentence, return_tensors="pt", truncation=True, max_length=512)
+        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()}
@@ -64,179 +74,255 @@ def translate(text, source_lang, target_lang):
     return '\n'.join(translated_sentences)
 
 def swap_languages(src, tgt, input_txt, output_txt):
-    """交换语言方向"""
     return tgt, src, output_txt, input_txt
 
-# ==================== 示例数据 ====================
-EXAMPLES = {
-    "English": {
-        "drilling": "Mud weight adjusted to 1.82 specific gravity at 3,247 meters depth.",
-        "reservoir": "Permeability is 250 millidarcy with 22 percent porosity.",
-        "subsea": "Christmas tree rated for 10,000 psi working pressure.",
-        "safety": "H2S training required before site access."
-    },
-    "Norwegian": {
-        "drilling": "Slamvekt justert til 1,82 spesifikk tyngde ved 3 247 meters dybde.",
-        "reservoir": "Permeabilitet er 250 millidarcy med 22 prosent porøsitet.",
-        "subsea": "Juletre dimensjonert for 10 000 psi arbeidstrykk.",
-        "safety": "H2S-opplæring påkrevd før tilgang til området."
-    }
+def load_file(file):
+    if file is None:
+        return ""
+    
+    try:
+        with open(file.name, 'r', encoding='utf-8') as f:
+            return f.read()
+    except:
+        try:
+            with open(file.name, 'r', encoding='latin-1') as f:
+                return f.read()
+        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."
 }
 
-# ==================== CSS 样式 ====================
+# 挪威语示例
+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
+
+def update_example_buttons(source_lang):
+    """当语言改变时更新示例"""
+    examples = get_examples(source_lang)
+    return examples["drilling_short"]
+
 custom_css = """
 .gradio-container {
-    max-width: 1200px !important;
-    margin: 0 auto !important;
+    max-width: 1100px !important;
+    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
+}
+.main-container {
+    background: #f6f7f8 !important;
+    padding: 0 !important;
+    border-radius: 0 !important;
 }
 .translate-box {
-    background: white;
-    border: 1px solid #e5e7eb;
-    border-radius: 8px;
-    padding: 0;
-    box-shadow: 0 1px 3px rgba(0,0,0,0.1);
+    background: white !important;
+    border-radius: 5px !important;
+    padding: 0 !important;
+    box-shadow: 0 2px 4px rgba(0,0,0,0.08) !important;
+    margin: 20px 0 !important;
+}
+.lang-header {
+    padding: 16px 20px !important;
+    border-bottom: 1px solid #e8eaed !important;
+    background: #fafafa !important;
 }
 .lang-selector {
-    font-size: 16px !important;
-    font-weight: 600 !important;
-    padding: 12px !important;
-    background: #f9fafb !important;
-    border-radius: 8px 8px 0 0 !important;
+    border: none !important;
+    background: transparent !important;
+    font-size: 15px !important;
+    font-weight: 500 !important;
+    color: #333 !important;
 }
 .text-area textarea {
-    font-size: 15px !important;
-    line-height: 1.6 !important;
-    padding: 16px !important;
-    min-height: 180px !important;
     border: none !important;
-    resize: vertical !important;
+    font-size: 17px !important;
+    line-height: 1.7 !important;
+    padding: 20px !important;
+    min-height: 200px !important;
+}
+.swap-container {
+    display: flex !important;
+    align-items: center !important;
+    justify-content: center !important;
+    padding: 20px 0 !important;
 }
 .swap-btn {
-    width: 48px !important;
-    height: 48px !important;
+    width: 44px !important;
+    height: 44px !important;
+    min-width: 44px !important;
     border-radius: 50% !important;
-    font-size: 20px !important;
     background: white !important;
-    border: 2px solid #0ea5e9 !important;
-    color: #0ea5e9 !important;
-    transition: all 0.2s !important;
+    border: 1px solid #d1d5db !important;
+    box-shadow: 0 1px 3px rgba(0,0,0,0.1) !important;
+    font-size: 18px !important;
+    color: #0f6fff !important;
+    cursor: pointer !important;
 }
 .swap-btn:hover {
-    background: #0ea5e9 !important;
-    color: white !important;
-    transform: rotate(180deg) !important;
+    background: #f8f9fa !important;
+    border-color: #0f6fff !important;
 }
-.example-btn {
-    padding: 8px 16px !important;
-    font-size: 14px !important;
-    border-radius: 6px !important;
-    background: #f1f5f9 !important;
-    border: 1px solid #cbd5e1 !important;
-    transition: all 0.2s !important;
-}
-.example-btn:hover {
-    background: #e2e8f0 !important;
-    border-color: #94a3b8 !important;
+.footer-info {
+    text-align: center !important;
+    color: #999 !important;
+    font-size: 13px !important;
+    padding: 20px !important;
 }
 """
 
-# ==================== Gradio 界面 ====================
-with gr.Blocks(css=custom_css, theme=gr.themes.Soft(), title="Oil & Gas Translator") as demo:
+with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
     
-    gr.Markdown(
-        """
-        # 🛢️ Oil & Gas Technical Translator
-        ### English ↔ Norwegian Bidirectional Translation
-        """
-    )
+    gr.HTML("<div style='height: 20px'></div>")
     
-    # 主翻译区域
     with gr.Row():
-        with gr.Column(scale=5):
+        with gr.Column(scale=1):
             with gr.Group(elem_classes="translate-box"):
-                source_lang = gr.Dropdown(
-                    choices=["English", "Norwegian"],
-                    value="English",
-                    label="Source Language",
-                    elem_classes="lang-selector"
-                )
+                with gr.Row(elem_classes="lang-header"):
+                    source_lang = gr.Dropdown(
+                        choices=["English", "Norwegian"],
+                        value="English",
+                        show_label=False,
+                        container=False,
+                        elem_classes="lang-selector",
+                        scale=1
+                    )
+                
                 input_text = gr.Textbox(
-                    placeholder="Enter text to translate...",
+                    placeholder="Type to translate",
                     show_label=False,
-                    lines=6,
+                    lines=8,
+                    max_lines=20,
+                    container=False,
                     elem_classes="text-area"
                 )
         
-        with gr.Column(scale=1, min_width=80):
-            gr.HTML("<div style='height: 60px'></div>")
-            swap_btn = gr.Button("⇄", elem_classes="swap-btn", size="sm")
+        with gr.Column(scale=0, min_width=100):
+            with gr.Row(elem_classes="swap-container"):
+                swap_btn = gr.Button("⇄", elem_classes="swap-btn")
         
-        with gr.Column(scale=5):
+        with gr.Column(scale=1):
             with gr.Group(elem_classes="translate-box"):
-                target_lang = gr.Dropdown(
-                    choices=["English", "Norwegian"],
-                    value="Norwegian",
-                    label="Target Language",
-                    elem_classes="lang-selector"
-                )
+                with gr.Row(elem_classes="lang-header"):
+                    target_lang = gr.Dropdown(
+                        choices=["English", "Norwegian"],
+                        value="Norwegian",
+                        show_label=False,
+                        container=False,
+                        elem_classes="lang-selector",
+                        scale=1
+                    )
+                
                 output_text = gr.Textbox(
-                    placeholder="Translation will appear here...",
+                    placeholder="Translation",
                     show_label=False,
-                    lines=6,
+                    lines=8,
+                    max_lines=20,
+                    container=False,
                     elem_classes="text-area",
                     interactive=False
                 )
     
-    gr.Markdown("---")
+    gr.HTML(
+        "<div class='footer-info'>"
+        "Oil & Gas Translation • English ↔ Norwegian • Bidirectional Model"
+        "</div>"
+    )
     
-    # 示例区域
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("### 📝 Quick Examples")
-            gr.Markdown("Click an example below to try it:")
-            
-            with gr.Row():
-                drilling_btn = gr.Button("⛏️ Drilling", elem_classes="example-btn", scale=1)
-                reservoir_btn = gr.Button("🗂️ Reservoir", elem_classes="example-btn", scale=1)
-                subsea_btn = gr.Button("🌊 Subsea", elem_classes="example-btn", scale=1)
-                safety_btn = gr.Button("⚠️ Safety", elem_classes="example-btn", scale=1)
+    with gr.Accordion("Example Sentences", open=True):
+        with gr.Row():
+            example_text = gr.Textbox(
+                value=EXAMPLES_EN["drilling_short"],
+                label="",
+                lines=3,
+                max_lines=5,
+                show_copy_button=True
+            )
+            use_example_btn = gr.Button("Use This Example →", variant="primary", size="sm")
+        
+        with gr.Row():
+            btn1 = gr.Button("Drilling (Short)", size="sm")
+            btn2 = gr.Button("Drilling (Long)", size="sm")
+            btn3 = gr.Button("Reservoir (Short)", size="sm")
+            btn4 = gr.Button("Reservoir (Long)", size="sm")
+            btn5 = gr.Button("Subsea (Short)", size="sm")
+        
+        with gr.Row():
+            btn6 = gr.Button("Subsea (Long)", size="sm")
+            btn7 = gr.Button("Seismic (Short)", size="sm")
+            btn8 = gr.Button("Seismic (Long)", size="sm")
+            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)
+        btn4.click(lambda sl: get_example("reservoir_long", sl), inputs=[source_lang], outputs=example_text)
+        btn5.click(lambda sl: get_example("subsea_short", sl), inputs=[source_lang], outputs=example_text)
+        btn6.click(lambda sl: get_example("subsea_long", sl), inputs=[source_lang], outputs=example_text)
+        btn7.click(lambda sl: get_example("seismic_short", sl), inputs=[source_lang], outputs=example_text)
+        btn8.click(lambda sl: get_example("seismic_long", sl), inputs=[source_lang], outputs=example_text)
+        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
+        )
     
-    # 文件上传
-    gr.Markdown("---")
-    with gr.Accordion("📄 Upload Text File", open=False):
+    with gr.Accordion("Upload Text File", open=False):
         file_input = gr.File(
             label="Upload a .txt file to translate",
-            file_types=[".txt"]
+            file_types=[".txt"],
+            type="filepath"
         )
     
-    # ==================== 事件绑定 ====================
-    
-    # 示例按钮点击
-    def load_example(example_key, source_lang):
-        return EXAMPLES[source_lang][example_key]
-    
-    drilling_btn.click(
-        fn=lambda sl: load_example("drilling", sl),
-        inputs=[source_lang],
-        outputs=[input_text]
-    )
-    reservoir_btn.click(
-        fn=lambda sl: load_example("reservoir", sl),
-        inputs=[source_lang],
-        outputs=[input_text]
-    )
-    subsea_btn.click(
-        fn=lambda sl: load_example("subsea", sl),
-        inputs=[source_lang],
-        outputs=[input_text]
-    )
-    safety_btn.click(
-        fn=lambda sl: load_example("safety", sl),
+    # 当源语言改变时,更新示例显示
+    source_lang.change(
+        fn=update_example_buttons,
         inputs=[source_lang],
-        outputs=[input_text]
+        outputs=[example_text]
     )
     
-    # 翻译触发
     input_text.change(
         fn=translate,
         inputs=[input_text, source_lang, target_lang],
@@ -255,41 +341,16 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft(), title="Oil & Gas Translat
         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]
     )
     
-    # 文件上传
-    def load_file(file):
-        if file is None:
-            return ""
-        try:
-            with open(file.name, 'r', encoding='utf-8') as f:
-                return f.read()
-        except:
-            try:
-                with open(file.name, 'r', encoding='latin-1') as f:
-                    return f.read()
-            except Exception as e:
-                return f"Error: {str(e)}"
-    
     file_input.change(
         fn=load_file,
         inputs=file_input,
         outputs=input_text
     )
-    
-    # 页脚
-    gr.Markdown(
-        """
-        ---
-        <div style='text-align: center; color: #64748b; font-size: 13px;'>
-        Powered by NLLB-200 + LoRA | Specialized for Oil & Gas Industry
-        </div>
-        """
-    )
 
-demo.launch(share=False)
+demo.launch()