app.py

import gradio as gr
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(
    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)

def translate(text, source_lang, target_lang):
    if not text.strip():
        return ""
    
    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 = []
    
    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)
    
    return '\n'.join(translated_sentences)

def swap_languages(src, tgt, input_txt, output_txt):
    return tgt, src, output_txt, input_txt

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."
}

# 挪威语示例
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: 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 !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 {
    border: none !important;
    background: transparent !important;
    font-size: 15px !important;
    font-weight: 500 !important;
    color: #333 !important;
}
.text-area textarea {
    border: none !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: 44px !important;
    height: 44px !important;
    min-width: 44px !important;
    border-radius: 50% !important;
    background: white !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: #f8f9fa !important;
    border-color: #0f6fff !important;
}
.footer-info {
    text-align: center !important;
    color: #999 !important;
    font-size: 13px !important;
    padding: 20px !important;
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
    
    gr.HTML("<div style='height: 20px'></div>")
    
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Group(elem_classes="translate-box"):
                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="Type to translate",
                    show_label=False,
                    lines=8,
                    max_lines=20,
                    container=False,
                    elem_classes="text-area"
                )
        
        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=1):
            with gr.Group(elem_classes="translate-box"):
                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",
                    show_label=False,
                    lines=8,
                    max_lines=20,
                    container=False,
                    elem_classes="text-area",
                    interactive=False
                )
    
    gr.HTML(
        "<div class='footer-info'>"
        "Oil & Gas Translation • English ↔ Norwegian • Bidirectional Model"
        "</div>"
    )
    
    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
        )
    
    with gr.Accordion("Upload Text File", open=False):
        file_input = gr.File(
            label="Upload a .txt file to translate",
            file_types=[".txt"],
            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
    )

demo.launch()