Update app.py

app.py

@@ -1,75 +1,95 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# --- MODEL DATA ---
-MODELS = {
+# --- CONFIGURATION ---
+MODELS_CONFIG = {
     "Phase 2: Stable (Formal)": {
         "id": "st192011/Maltese-EuroLLM-1.7B-Phase2-Stable",
-        "description": (
-            "The 'Bureaucrat Bot'. Built upon a foundational adaptation phase that mixed "
-            "monolingual Maltese and Italian to bridge morphological roots. This version "
-            "was fine-tuned on high-fidelity EU and governmental parallel corpora, "
-            "optimizing it for extreme formal precision and administrative accuracy."
-        ),
+        "description": "The 'Bureaucrat Bot'. Optimized for formal precision.",
         "chrf": "60.18",
         "comet": "0.6431"
     },
     "Phase 4: Anchored (Native)": {
         "id": "st192011/Maltese-EuroLLM-1.7B-Phase4-Anchored",
-        "description": (
-            "The 'Native Speaker'. An evolution of Phase 2 utilizing a curriculum-based "
-            "'Full Circle' approach. It integrates synthesized reasoning chains (CoT) "
-            "that allow the model to process linguistic logic before translating. By mixing "
-            "all previous data types, it anchors factual accuracy to native-level phrasing "
-            "and cultural awareness."
-        ),
+        "description": "The 'Native Speaker'. Optimized for cultural awareness and logic.",
         "chrf": "52.68",
         "comet": "0.6567"
     }
 }
 
+# --- MODEL LOADING ---
+# We load them globally so they stay in memory (this requires ~14GB RAM total)
+print("Loading models to CPU... this may take a few minutes.")
+
+# Load Model 2
+tokenizer_p2 = AutoTokenizer.from_pretrained(MODELS_CONFIG["Phase 2: Stable (Formal)"]["id"])
+model_p2 = AutoModelForCausalLM.from_pretrained(
+    MODELS_CONFIG["Phase 2: Stable (Formal)"]["id"], 
+    device_map="cpu", 
+    torch_dtype=torch.float32
+)
+
+# Load Model 4
+tokenizer_p4 = AutoTokenizer.from_pretrained(MODELS_CONFIG["Phase 4: Anchored (Native)"]["id"])
+model_p4 = AutoModelForCausalLM.from_pretrained(
+    MODELS_CONFIG["Phase 4: Anchored (Native)"]["id"], 
+    device_map="cpu", 
+    torch_dtype=torch.float32
+)
+
+def local_translate(model, tokenizer, text, temp):
+    prompt = f"### INGLIŻ: {text}\n### MALTI:"
+    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
+    
+    with torch.no_grad():
+        output_tokens = model.generate(
+            **inputs,
+            max_new_tokens=150,
+            temperature=temp,
+            do_sample=True if temp > 0.1 else False,
+            repetition_penalty=1.2,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    # Decode only the new tokens
+    full_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+    # Extract the part after ### MALTI:
+    maltese_text = full_text.split("### MALTI:")[-1].strip()
+    return maltese_text
+
 def translate_logic(text, selected_models, temp):
-    results = {}
+    out_p2 = "Model not selected."
+    out_p4 = "Model not selected."
     
-    for model_name in selected_models:
-        model_id = MODELS[model_name]["id"]
-        client = InferenceClient(model=model_id)
-        
-        # Prompt format consistent with training
-        prompt = f"### INGLIŻ: {text}\n### MALTI:"
-        
+    if not text.strip():
+        return "Please enter text.", "Please enter text."
+
+    if "Phase 2: Stable (Formal)" in selected_models:
         try:
-            output = client.text_generation(
-                prompt,
-                max_new_tokens=150,
-                temperature=temp,
-                do_sample=True if temp > 0.1 else False,
-                repetition_penalty=1.2
-            )
-            # Clean up the response
-            clean_output = output.strip().split("### MALTI:")[-1].replace("<|endoftext|>", "").strip()
-            results[model_name] = clean_output
+            out_p2 = local_translate(model_p2, tokenizer_p2, text, temp)
+        except Exception as e:
+            out_p2 = f"Error: {str(e)}"
+
+    if "Phase 4: Anchored (Native)" in selected_models:
+        try:
+            out_p4 = local_translate(model_p4, tokenizer_p4, text, temp)
         except Exception as e:
-            results[model_name] = f"Error: Inference API is still loading or unavailable. ({str(e)})"
+            out_p4 = f"Error: {str(e)}"
             
-    # Return formatted outputs for the UI
-    # We return a list of outputs corresponding to the two textboxes
-    out_p2 = results.get("Phase 2: Stable (Formal)", "Model not selected.")
-    out_p4 = results.get("Phase 4: Anchored (Native)", "Model not selected.")
-    
     return out_p2, out_p4
 
 # --- GRADIO UI ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🇲🇹 Maltese-MT Lab")
-    gr.Markdown("Compare different generations of fine-tuned EuroLLM models for English-to-Maltese translation.")
+    gr.Markdown("# 🇲🇹 Maltese-MT Lab (Local CPU)")
+    gr.Markdown("Comparing English-to-Maltese EuroLLM models running directly on this machine.")
     
     with gr.Row():
         with gr.Column(scale=2):
             input_text = gr.Textbox(label="English Source Text", placeholder="Enter English text here...", lines=4)
             model_selector = gr.CheckboxGroup(
-                choices=list(MODELS.keys()), 
-                value=list(MODELS.keys()), 
+                choices=list(MODELS_CONFIG.keys()), 
+                value=list(MODELS_CONFIG.keys()), 
                 label="Select Models to Compare"
             )
             temp_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.1, step=0.1, label="Creativity (Temperature)")
@@ -79,14 +99,12 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         with gr.Column():
             gr.Markdown("### Phase 2: Stable")
             p2_out = gr.Textbox(label="Output", interactive=False, lines=5)
-            gr.Markdown(f"**Training:** {MODELS['Phase 2: Stable (Formal)']['description']}")
-            gr.Markdown(f"**Metrics:** ChrF++: `{MODELS['Phase 2: Stable (Formal)']['chrf']}` | COMET: `{MODELS['Phase 2: Stable (Formal)']['comet']}`")
+            gr.Markdown(f"**ChrF++:** `{MODELS_CONFIG['Phase 2: Stable (Formal)']['chrf']}` | **COMET:** `{MODELS_CONFIG['Phase 2: Stable (Formal)']['comet']}`")
             
         with gr.Column():
             gr.Markdown("### Phase 4: Anchored")
             p4_out = gr.Textbox(label="Output", interactive=False, lines=5)
-            gr.Markdown(f"**Training:** {MODELS['Phase 4: Anchored (Native)']['description']}")
-            gr.Markdown(f"**Metrics:** ChrF++: `{MODELS['Phase 4: Anchored (Native)']['chrf']}` | COMET: `{MODELS['Phase 4: Anchored (Native)']['comet']}`")
+            gr.Markdown(f"**ChrF++:** `{MODELS_CONFIG['Phase 4: Anchored (Native)']['chrf']}` | **COMET:** `{MODELS_CONFIG['Phase 4: Anchored (Native)']['comet']}`")
 
     gr.Examples(
         examples=[