add fine-tuned

app.py

@@ -1,18 +1,26 @@
 import gradio as gr
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import SentenceTransformer, CrossEncoder
 from supabase import create_client
 import os
 from dotenv import load_dotenv
 
 load_dotenv()
 
-model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+embedder = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
 
 # === Supabase ===
 supabase_url = os.getenv("SUPABASE_URL")
 supabase_key = os.getenv("SUPABASE_KEY")
 supabase = create_client(supabase_url, supabase_key)
 
+def expand_query(query: str) -> str:
+    """
+    Fungsi kecil untuk memperluas query.
+    Bisa ditambah sinonim, kata kunci, atau parafrasa.
+    """
+    return f"Kegiatan usaha yang berkaitan dengan {query}"
+
 def get_embedding(text: str):
     """
     Menghasilkan embedding vector dari teks menggunakan model SentenceTransformer.
@@ -20,9 +28,32 @@ def get_embedding(text: str):
     if not text:
         return []
     
-    embedding = model.encode(text).tolist()
+    expanded_text = expand_query(text)
+    embedding = embedder.encode(expanded_text).tolist()
     return embedding
 
+def fn_semantic(query: str, match_count: int = 25):
+    embedding = embedder.encode(query).tolist()
+
+    response = supabase.rpc(
+        "search_kbli",
+        {"query_embedding": embedding, "match_count": match_count}
+    ).execute()
+    candidates = response.data or []
+
+    if not candidates:
+        return {"results": []}
+
+    pairs = [(query, c["judul"] + " " + c["deskripsi"]) for c in candidates]
+    scores = reranker.predict(pairs)
+
+    for c, s in zip(candidates, scores):
+        c["rerank_score"] = float(s)
+
+    candidates = sorted(candidates, key=lambda x: x["rerank_score"], reverse=True)
+
+    return {"results": candidates[:10]}
+
 def search_kbli(text: str):
     if not text:
         return {"embedding": [], "results": []}
@@ -156,19 +187,35 @@ with gr.Blocks(css="""
         inp1.submit(get_embedding, inp1, out1, api_name="get_embedding")
         btn_clear1.click(lambda: ("", None), None, [inp1, out1])
         btn_submit1.click(get_embedding, inp1, out1, api_name="get_embedding")
+    
+    with gr.Tab("Embedding Fine-tuned"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                inp2 = gr.Textbox(label="Masukkan teks")
+
+                with gr.Row(elem_classes="btn-row"):
+                    btn_clear2 = gr.Button("Clear", variant="secondary")
+                    btn_submit2 = gr.Button("Submit", variant="primary")
+
+            with gr.Column(scale=1):
+                out2 = gr.JSON(label="Embedding Vector")
+        
+        inp2.submit(fn_semantic, inp2, out2, api_name="fn_semantic")
+        btn_clear2.click(lambda: ("", None), None, [inp2, out2])
+        btn_submit2.click(fn_semantic, inp2, out2, api_name="fn_semantic")
 
     with gr.Tab("Search KBLI"):
-        inp2 = gr.Textbox(label="Masukkan teks")
+        inp3 = gr.Textbox(label="Masukkan teks")
 
         with gr.Row(elem_classes="btn-row-search"):
-            btn_clear2 = gr.Button("Clear", variant="secondary")
-            btn_submit2 = gr.Button("Submit", variant="primary")
+            btn_clear3 = gr.Button("Clear", variant="secondary")
+            btn_submit3 = gr.Button("Submit", variant="primary")
         
-        out2 = gr.HTML(label="Hasil Pencarian Semantic")
+        out3 = gr.HTML(label="Hasil Pencarian Semantic")
         
-        inp2.submit(search_kbli, inp2, out2, api_name="search_kbli")
-        btn_clear2.click(lambda: ("", None), None, [inp2, out2])
-        btn_submit2.click(search_kbli, inp2, out2, api_name="search_kbli")
+        inp3.submit(search_kbli, inp3, out3, api_name="search_kbli")
+        btn_clear3.click(lambda: ("", None), None, [inp3, out3])
+        btn_submit3.click(search_kbli, inp3, out3, api_name="search_kbli")
 
 if __name__ == "__main__":
     demo.queue().launch()