Update app.py

app.py

@@ -1,7 +1,9 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, BitsAndBytesConfig
 from peft import PeftModel
+from functools import lru_cache
+import os
 
 base_model_name = "facebook/nllb-200-distilled-600M"
 adapter_en_to_no = "entropy25/mt_en_no_oil"
@@ -9,58 +11,122 @@ adapter_no_to_en = "entropy25/mt_no_en_oil"
 
 tokenizer = AutoTokenizer.from_pretrained(base_model_name)
 
-print("Loading shared base model...")
+print("Loading shared base model with 8-bit quantization...")
+quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+
 base_model = AutoModelForSeq2SeqLM.from_pretrained(
     base_model_name,
-    torch_dtype=torch.float16,
-    low_cpu_mem_usage=True,
-    device_map="auto"
+    quantization_config=quantization_config,
+    device_map="auto",
+    low_cpu_mem_usage=True
 )
 
 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")
+model.eval()
+
+QUALITY_PRESETS = {
+    "Professional (Best Quality)": {"num_beams": 3, "max_length": 256, "batch_size": 4},
+    "Balanced (Faster)": {"num_beams": 2, "max_length": 256, "batch_size": 5},
+    "Draft (Fastest)": {"num_beams": 2, "max_length": 128, "batch_size": 5}
+}
 
-def translate(text, source_lang, target_lang):
+QUALITY_TEST_CASES = {
+    "en_to_no": [
+        {
+            "input": "Mud weight adjusted to 1.82 specific gravity at 3,247 meters depth.",
+            "expected": "Slamvekt justert til 1,82 spesifikk tyngde ved 3 247 meters dybde.",
+            "check": ["slamvekt", "1,82", "3 247"]
+        },
+        {
+            "input": "Christmas tree rated for 10,000 psi working pressure.",
+            "expected": "Juletre dimensjonert for 10 000 psi arbeidstrykk.",
+            "check": ["juletre", "10 000", "psi"]
+        },
+        {
+            "input": "H2S training required before site access.",
+            "expected": "H2S-opplæring påkrevd før tilgang til området.",
+            "check": ["H2S", "opplæring", "påkrevd"]
+        },
+        {
+            "input": "Permeability is 250 millidarcy with 22 percent porosity.",
+            "expected": "Permeabilitet er 250 millidarcy med 22 prosent porøsitet.",
+            "check": ["permeabilitet", "250", "22"]
+        }
+    ],
+    "no_to_en": [
+        {
+            "input": "Permeabilitet er 250 millidarcy med 22 prosent porøsitet.",
+            "expected": "Permeability is 250 millidarcy with 22 percent porosity.",
+            "check": ["permeability", "250", "22"]
+        },
+        {
+            "input": "Subsea produksjonssystemet består av et vertikalt juletre.",
+            "expected": "The subsea production system consists of a vertical Christmas tree.",
+            "check": ["subsea", "Christmas tree", "vertical"]
+        },
+        {
+            "input": "Slamvekt justert til 1,82 spesifikk tyngde ved 3 247 meters dybde.",
+            "expected": "Mud weight adjusted to 1.82 specific gravity at 3,247 meters depth.",
+            "check": ["mud weight", "1.82", "3,247"]
+        }
+    ]
+}
+
+MAX_FILE_SIZE = 1024 * 1024
+
+def translate_core(text, source_lang, target_lang, quality_preset):
     if not text.strip() or source_lang == target_lang:
         return text
     
     if source_lang == "English" and target_lang == "Norwegian":
         model.set_adapter("en_to_no")
-        src_code, tgt_code = "eng_Latn", "nob_Latn"
+        tgt_code = "nob_Latn"
     elif source_lang == "Norwegian" and target_lang == "English":
         model.set_adapter("no_to_en")
-        src_code, tgt_code = "nob_Latn", "eng_Latn"
+        tgt_code = "eng_Latn"
     else:
         return "Unsupported language pair"
     
+    preset = QUALITY_PRESETS[quality_preset]
+    
     lines = text.split('\n')
     non_empty_lines = [line for line in lines if line.strip()]
     
     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
-    )
+    batch_size = preset["batch_size"]
+    all_results = []
     
-    results = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    for i in range(0, len(non_empty_lines), batch_size):
+        batch = non_empty_lines[i:i+batch_size]
+        
+        inputs = tokenizer(
+            batch,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=preset["max_length"]
+        )
+        
+        if hasattr(model, 'device'):
+            inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        
+        with torch.inference_mode():
+            outputs = model.generate(
+                **inputs,
+                forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_code),
+                max_length=preset["max_length"],
+                num_beams=preset["num_beams"],
+                early_stopping=True
+            )
+        
+        batch_results = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        all_results.extend(batch_results)
     
-    result_iter = iter(results)
+    result_iter = iter(all_results)
     final_lines = []
     for line in lines:
         if line.strip():
@@ -70,19 +136,84 @@ def translate(text, source_lang, target_lang):
     
     return '\n'.join(final_lines)
 
+@lru_cache(maxsize=512)
+def translate_cached(text, source_lang, target_lang, quality_preset):
+    return translate_core(text, source_lang, target_lang, quality_preset)
+
+def translate(text, source_lang, target_lang, quality_preset):
+    if len(text) > 10000:
+        return "Error: Text too long (max 10,000 characters)"
+    return translate_cached(text, source_lang, target_lang, quality_preset)
+
+def run_quality_tests():
+    results = []
+    results.append("=== QUALITY REGRESSION TEST ===\n")
+    
+    for direction, test_cases in QUALITY_TEST_CASES.items():
+        if direction == "en_to_no":
+            src_lang, tgt_lang = "English", "Norwegian"
+        else:
+            src_lang, tgt_lang = "Norwegian", "English"
+        
+        results.append(f"\n{src_lang} to {tgt_lang}\n")
+        
+        for i, case in enumerate(test_cases, 1):
+            translation = translate_core(case["input"], src_lang, tgt_lang, "Professional (Best Quality)")
+            
+            passed_checks = []
+            failed_checks = []
+            
+            for keyword in case["check"]:
+                if keyword.lower() in translation.lower():
+                    passed_checks.append(keyword)
+                else:
+                    failed_checks.append(keyword)
+            
+            status = "PASS" if not failed_checks else "CHECK"
+            
+            results.append(f"\nTest {i}: {status}")
+            results.append(f"Input:    {case['input']}")
+            results.append(f"Expected: {case['expected']}")
+            results.append(f"Got:      {translation}")
+            
+            if passed_checks:
+                results.append(f"Found:  {', '.join(passed_checks)}")
+            if failed_checks:
+                results.append(f"Missing: {', '.join(failed_checks)}")
+    
+    results.append("\n=== TEST COMPLETE ===")
+    
+    pass_count = sum(1 for r in results if "PASS" in r)
+    check_count = sum(1 for r in results if "CHECK" in r)
+    total = len(QUALITY_TEST_CASES["en_to_no"]) + len(QUALITY_TEST_CASES["no_to_en"])
+    
+    results.insert(1, f"\nScore: {pass_count}/{total} passed, {check_count}/{total} need review\n")
+    
+    return '\n'.join(results)
+
 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:
+        if os.path.getsize(file.name) > MAX_FILE_SIZE:
+            return "Error: File too large (max 1MB)"
+        
         with open(file.name, 'r', encoding='utf-8') as f:
-            return f.read()
+            content = f.read()
+            if len(content) > 10000:
+                return "Error: File content too long (max 10,000 characters)"
+            return content
     except:
         try:
             with open(file.name, 'r', encoding='latin-1') as f:
-                return f.read()
+                content = f.read()
+                if len(content) > 10000:
+                    return "Error: File content too long (max 10,000 characters)"
+                return content
         except Exception as e:
             return f"Error reading file: {str(e)}"
 
@@ -181,17 +312,44 @@ custom_css = """
     background: #f8f9fa !important;
     border-color: #0f6fff !important;
 }
+.translate-btn {
+    background: #0f6fff !important;
+    color: white !important;
+    border: none !important;
+    padding: 12px 24px !important;
+    font-size: 15px !important;
+    font-weight: 500 !important;
+    border-radius: 4px !important;
+    cursor: pointer !important;
+}
+.translate-btn:hover {
+    background: #0d5dd9 !important;
+}
 .footer-info {
     text-align: center !important;
     color: #999 !important;
     font-size: 13px !important;
     padding: 20px !important;
 }
+.quality-selector {
+    background: #f0f7ff !important;
+    border: 1px solid #0f6fff !important;
+    border-radius: 4px !important;
+}
 """
 
 with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
     gr.HTML("<div style='height: 20px'></div>")
     
+    with gr.Row():
+        quality_preset = gr.Radio(
+            choices=list(QUALITY_PRESETS.keys()),
+            value="Professional (Best Quality)",
+            label="Translation Quality",
+            info="Professional: beam=3, max=256 | Balanced: beam=2, max=256 | Draft: beam=2, max=128",
+            elem_classes="quality-selector"
+        )
+    
     with gr.Row():
         with gr.Column(scale=1):
             with gr.Group(elem_classes="translate-box"):
@@ -238,7 +396,10 @@ 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>")
+    with gr.Row():
+        translate_btn = gr.Button("Translate", variant="primary", elem_classes="translate-btn", size="lg")
+    
+    gr.HTML("<div class='footer-info'>Oil & Gas Translation • English ↔ Norwegian • Optimized for HF Space</div>")
     
     with gr.Accordion("Example Sentences", open=True):
         with gr.Row():
@@ -249,7 +410,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
                 max_lines=5,
                 show_copy_button=True
             )
-            use_example_btn = gr.Button("Use This Example →", variant="primary", size="sm")
+            use_example_btn = gr.Button("Use This Example", variant="primary", size="sm")
         
         with gr.Row():
             btn1 = gr.Button("Drilling (Short)", size="sm")
@@ -280,16 +441,35 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:
     
     with gr.Accordion("Upload Text File", open=False):
         file_input = gr.File(
-            label="Upload a .txt file to translate",
+            label="Upload a .txt file to translate (max 1MB)",
             file_types=[".txt"],
             type="filepath"
         )
     
+    with gr.Accordion("Quality Test (Developer)", open=False):
+        test_output = gr.Textbox(
+            label="Test Results",
+            lines=20,
+            max_lines=30,
+            interactive=False
+        )
+        run_test_btn = gr.Button("Run Quality Regression Test", variant="secondary")
+        run_test_btn.click(fn=run_quality_tests, outputs=test_output)
+    
     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])
+    
+    translate_btn.click(
+        fn=translate,
+        inputs=[input_text, source_lang, target_lang, quality_preset],
+        outputs=output_text
+    )
+    
+    swap_btn.click(
+        fn=swap_languages,
+        inputs=[source_lang, target_lang, input_text, output_txt],
+        outputs=[source_lang, target_lang, input_text, output_text]
+    )
+    
     file_input.change(fn=load_file, inputs=file_input, outputs=input_text)
 
-demo.launch()
+demo.queue(concurrency_count=1, max_size=20).launch()