mt / app.py
entropy25's picture
Update app.py
56b2235 verified
import gradio as gr
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"
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
base_model = AutoModelForSeq2SeqLM.from_pretrained(
base_model_name,
torch_dtype=torch.float16,
low_cpu_mem_usage=True,
device_map="auto"
)
model = PeftModel.from_pretrained(base_model, adapter_model_name)
def translate(text, source_lang, target_lang):
if not text.strip():
return ""
if source_lang == target_lang:
return text
lang_map = {
"English": "eng_Latn",
"Norwegian": "nob_Latn"
}
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(lang_map[target_lang]),
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)}"
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;
}
"""
EXAMPLES = {
"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": "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": "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": "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 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."
}
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"
"</div>"
)
with gr.Accordion("Example Sentences", open=True):
with gr.Row():
example_text = gr.Textbox(
value=EXAMPLES["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():
gr.Button("Drilling (Short)", size="sm").click(
lambda: EXAMPLES["drilling_short"],
outputs=example_text
)
gr.Button("Drilling (Long)", size="sm").click(
lambda: EXAMPLES["drilling_long"],
outputs=example_text
)
gr.Button("Reservoir (Short)", size="sm").click(
lambda: EXAMPLES["reservoir_short"],
outputs=example_text
)
gr.Button("Reservoir (Long)", size="sm").click(
lambda: EXAMPLES["reservoir_long"],
outputs=example_text
)
gr.Button("Subsea (Short)", size="sm").click(
lambda: EXAMPLES["subsea_short"],
outputs=example_text
)
with gr.Row():
gr.Button("Subsea (Long)", size="sm").click(
lambda: EXAMPLES["subsea_long"],
outputs=example_text
)
gr.Button("Seismic (Short)", size="sm").click(
lambda: EXAMPLES["seismic_short"],
outputs=example_text
)
gr.Button("Seismic (Long)", size="sm").click(
lambda: EXAMPLES["seismic_long"],
outputs=example_text
)
gr.Button("Safety (Short)", size="sm").click(
lambda: EXAMPLES["safety_short"],
outputs=example_text
)
gr.Button("Safety (Long)", size="sm").click(
lambda: EXAMPLES["safety_long"],
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"
)
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()