Update app.py

app.py

@@ -1,104 +1,89 @@
 # ============================================================
-# FINAL SYSTEM
-# NLLB + LLaMA + RAG
-# Javanese → Indonesian → English
+# FINAL TRANSLATION SYSTEM
+# NLLB + RAG + OPEN LLM (NO GATED MODEL)
 # ============================================================
 
-import os
 import torch
 import faiss
 import pandas as pd
 import gradio as gr
-import numpy as np
-
+from sentence_transformers import SentenceTransformer
 from transformers import (
     AutoTokenizer,
     AutoModelForSeq2SeqLM,
-    AutoModelForCausalLM,
+    AutoModelForCausalLM
 )
-from sentence_transformers import SentenceTransformer
 
 # ============================================================
-# Device
+# DEVICE
 # ============================================================
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # ============================================================
-# 1. LOAD NLLB (TRANSLATION ENGINE)
+# 1. NLLB TRANSLATION MODEL
 # ============================================================
 
 NLLB_MODEL = "facebook/nllb-200-distilled-600M"
 
 nllb_tokenizer = AutoTokenizer.from_pretrained(NLLB_MODEL)
-nllb_model = AutoModelForSeq2SeqLM.from_pretrained(NLLB_MODEL)
-nllb_model.to(device)
+nllb_model = AutoModelForSeq2SeqLM.from_pretrained(NLLB_MODEL).to(device)
 nllb_model.eval()
 
 JAV = "jav_Latn"
 IND = "ind_Latn"
-ENG = "eng_Latn"
 
 # ============================================================
-# 2. LOAD LLaMA (REFINEMENT ENGINE)
+# 2. OPEN LLM (REPLACEMENT FOR LLAMA)
 # ============================================================
 
-LLAMA_MODEL = "meta-llama/Llama-3.1-8B-Instruct"
+LLM_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"
+# alternatif:
+# "Qwen/Qwen2.5-7B-Instruct"
 
-llama_tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL)
-llama_model = AutoModelForCausalLM.from_pretrained(
-    LLAMA_MODEL,
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-    device_map="auto"
+llm_tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
+llm_model = AutoModelForCausalLM.from_pretrained(
+    LLM_MODEL,
+    device_map="auto",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
 )
-llama_model.eval()
+llm_model.eval()
 
 # ============================================================
 # 3. LOAD KNOWLEDGE BASE
 # ============================================================
 
-KB_PATH = "kb_jawa_ngoko_krama_indonesia_100k.csv"
-
-kb = pd.read_csv(KB_PATH)
+kb = pd.read_csv("kb_jawa_ngoko_krama_indonesia_100k.csv")
 kb["jv"] = kb["jv"].astype(str)
 kb["id"] = kb["id"].astype(str)
 
-kb_pairs = list(zip(kb["jv"], kb["id"]))
+pairs = list(zip(kb["jv"], kb["id"]))
 
 # ============================================================
-# 4. EMBEDDING + FAISS
+# 4. FAISS RAG
 # ============================================================
 
 embedder = SentenceTransformer(
     "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
 )
 
-jv_texts = kb["jv"].tolist()
-embeddings = embedder.encode(jv_texts, convert_to_numpy=True)
+emb = embedder.encode(kb["jv"].tolist(), convert_to_numpy=True)
 
-index = faiss.IndexFlatL2(embeddings.shape[1])
-index.add(embeddings)
+index = faiss.IndexFlatL2(emb.shape[1])
+index.add(emb)
 
-# ============================================================
-# 5. RETRIEVE PARALLEL PHRASES
-# ============================================================
-
-def retrieve_pairs(jv_text, k=5):
-    vec = embedder.encode([jv_text])
-    _, I = index.search(vec, k)
-
-    results = []
-    for i in I[0]:
-        jv, idn = kb_pairs[i]
-        results.append(f"- {jv} → {idn}")
-
-    return "\n".join(results)
+def retrieve(text, k=5):
+    v = embedder.encode([text])
+    _, I = index.search(v, k)
+    return "\n".join(
+        [f"- {pairs[i][0]} → {pairs[i][1]}" for i in I[0]]
+    )
 
 # ============================================================
-# 6. NLLB TRANSLATION
+# 5. NLLB TRANSLATE
 # ============================================================
 
-def nllb_translate(text, tgt_lang):
+def nllb_translate(text):
     nllb_tokenizer.src_lang = JAV
 
     inputs = nllb_tokenizer(
@@ -108,111 +93,71 @@ def nllb_translate(text, tgt_lang):
         max_length=512
     ).to(device)
 
-    output = nllb_model.generate(
+    out = nllb_model.generate(
         **inputs,
-        forced_bos_token_id=nllb_tokenizer.convert_tokens_to_ids(tgt_lang),
+        forced_bos_token_id=nllb_tokenizer.convert_tokens_to_ids(IND),
         max_length=256
     )
 
-    return nllb_tokenizer.batch_decode(
-        output, skip_special_tokens=True
-    )[0]
+    return nllb_tokenizer.decode(out[0], skip_special_tokens=True)
 
 # ============================================================
-# 7. LLaMA REFINEMENT
+# 6. LLM REFINEMENT
 # ============================================================
 
-def llama_refine(raw_translation, rag_context):
+def refine(raw, context):
 
     prompt = f"""
-Perbaiki terjemahan berikut agar alami dan gramatikal
-dalam Bahasa Indonesia.
+Perbaiki terjemahan berikut agar alami dalam Bahasa Indonesia.
 
-Terjemahan mesin:
-"{raw_translation}"
+Terjemahan awal:
+{raw}
 
-Referensi padanan bahasa:
-{rag_context}
+Referensi padanan:
+{context}
 
-Aturan:
-- Hanya keluarkan satu kalimat Bahasa Indonesia
-- Jangan menambahkan penjelasan
-- Jangan menyebut referensi
-- Jangan menambah informasi baru
+Instruksi:
+- Hasilkan satu kalimat Bahasa Indonesia
+- Jangan memberi penjelasan
 
 Jawaban:
 """
 
-    inputs = llama_tokenizer(
-        prompt,
-        return_tensors="pt"
-    ).to(llama_model.device)
+    inputs = llm_tokenizer(prompt, return_tensors="pt").to(llm_model.device)
 
-    output = llama_model.generate(
+    output = llm_model.generate(
         **inputs,
         max_new_tokens=80,
         temperature=0.2,
         do_sample=False
     )
 
-    text = llama_tokenizer.decode(
-        output[0],
-        skip_special_tokens=True
-    )
-
+    text = llm_tokenizer.decode(output[0], skip_special_tokens=True)
     return text.split("Jawaban:")[-1].strip()
 
 # ============================================================
-# 8. PIPELINE UTAMA
+# 7. PIPELINE
 # ============================================================
 
-def translate_pipeline(jv_text):
-
-    # Step 1 — Raw translation
-    raw_id = nllb_translate(jv_text, IND)
-
-    # Step 2 — RAG retrieval
-    rag_context = retrieve_pairs(jv_text)
-
-    # Step 3 — LLaMA refinement
-    final_id = llama_refine(raw_id, rag_context)
-
-    # Step 4 — English (optional)
-    final_en = nllb_translate(final_id, ENG)
-
-    return final_id, final_en
+def pipeline(jv_text):
+    raw = nllb_translate(jv_text)
+    ctx = retrieve(jv_text)
+    final = refine(raw, ctx)
+    return final
 
 # ============================================================
-# 9. GRADIO UI
+# 8. UI
 # ============================================================
 
-with gr.Blocks(title="NLLB + LLaMA + RAG Translator") as demo:
+with gr.Blocks() as demo:
+    gr.Markdown("## 🌾 NLLB + RAG + Open LLM Translator")
 
-    gr.Markdown("""
-    ## 🌾 Translator Bahasa Jawa  
-    **NLLB + LLaMA + RAG (Parallel Corpus)**
+    inp = gr.Textbox(label="Bahasa Jawa", lines=4)
+    out = gr.Textbox(label="Bahasa Indonesia", lines=4)
 
-    ✔ Low-resource language  
-    ✔ Linguistically grounded  
-    ✔ Research-grade architecture  
-    """)
+    btn = gr.Button("Terjemahkan")
+    btn.click(pipeline, inp, out)
 
-    inp = gr.Textbox(
-        label="Input Bahasa Jawa",
-        lines=4,
-        placeholder="Kula badhe sowan dhateng griya eyang."
-    )
-
-    out_id = gr.Textbox(label="Bahasa Indonesia", lines=4)
-    out_en = gr.Textbox(label="English", lines=4)
-
-    btn = gr.Button("🔄 Terjemahkan")
-    btn.click(translate_pipeline, inp, [out_id, out_en])
-
-demo.launch(
-    server_name="0.0.0.0",
-    server_port=7860,
-    ssr_mode=False
-)
+demo.launch()