refactor: update data loading and organization in main.py

.gitignore

@@ -1,8 +1,8 @@
 __pycache__/
 .gradio/
 cache/
-data/original/
 models/
 data/
 tmp/
+demo/
 .env

main.py

@@ -4,35 +4,40 @@ import pandas as pd
 import gradio as gr
 
 import faiss
+from datasets import load_dataset
 from sentence_transformers import SentenceTransformer
-from settings import OUTPUT_DIR, DEVICE
-os.environ['WANDB_DISABLED'] = 'true'
 
-    
-fine_tuned_model = SentenceTransformer(OUTPUT_DIR, device=DEVICE)
-passages         = pd.read_parquet('data/processed/corpus_data.parquet')['text'].tolist()
-legal_index      = faiss.read_index('data/retrieval/legal_faiss.index')
 
+# ===== Prepare model & data =====
+passages_path = hf_hub_download(repo_id='YuITC/Vietnamese-Legal-Doc-Retrieval-Data',
+                                filename='corpus_data.parquet', repo_type='dataset', 
+                                local_dir='demo')
+
+index_path = hf_hub_download(repo_id='YuITC/Vietnamese-Legal-Doc-Retrieval-Data', 
+                             filename='legal_faiss.index', repo_type='dataset', 
+                             local_dir='demo')
+
+emb_model   = SentenceTransformer('YuITC/bert-base-multilingual-cased-finetuned-VNLegalDocs')
+passages    = pd.read_parquet(passages_path)
+legal_index = faiss.read_index(index_path)
+
+
+# ===== Utility functions =====
 def retrieval(emb_model, query, index, top_k=10):
     q_emb = emb_model.encode(
         query, 
-        convert_to_numpy=True, 
-        normalize_embeddings=True,
+        convert_to_numpy=True, normalize_embeddings=True,
     ).astype(np.float32).reshape(1, -1)
     
-    scores, indices = index.search(q_emb, top_k)  # shape: (1, top_k)
-    
-    cand_idxs   = indices[0]
-    cand_scores = scores[0]
-    cand_texts  = [passages[i] for i in cand_idxs]
+    scores, indices = index.search(q_emb, top_k)
+    cand_idxs       = indices[0]
+    cand_scores     = scores[0]
+    cand_texts      = [passages[i] for i in cand_idxs]
 
-    results = [{
-        'index': int(cand_idxs[i]),
-        'score': float(cand_scores[i]),
-        'text': cand_texts[i]
-    } for i in range(len(cand_idxs))]
-    
-    return results
+    return [{'index': int(cand_idxs[i]),
+             'score': float(cand_scores[i]),
+             'text' : cand_texts[i]
+            } for i in range(len(cand_idxs))]
 
 def get_results(query, top_k):
     hits = retrieval(fine_tuned_model, query, legal_index, top_k=top_k)
@@ -43,8 +48,8 @@ def get_results(query, top_k):
     return result
 
     
+# ===== Gradio UI =====
 demo = gr.Interface(
-    'huggingface/YuITC/bert-base-multilingual-cased-finetuned-VNLegalDocs',
     fn=get_results,
     inputs=[
         gr.Textbox(lines=2, placeholder='Nhập câu hỏi pháp lý của bạn...', label='Câu hỏi'),