File size: 2,174 Bytes
0063d17 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import os
import numpy as np
import pandas as pd
import gradio as gr
import faiss
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')
def retrieval(emb_model, query, index, top_k=10):
q_emb = emb_model.encode(
query,
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]
results = [{
'index': int(cand_idxs[i]),
'score': float(cand_scores[i]),
'text': cand_texts[i]
} for i in range(len(cand_idxs))]
return results
def get_results(query, top_k):
hits = retrieval(fine_tuned_model, query, legal_index, top_k=top_k)
result = ""
for rank, h in enumerate(hits, start=1):
result += f"[Kết quả {rank} - Độ tin cậy={h['score']:.4f}]\n\n{h['text']}\n{'-'*100}\n"
return result
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'),
gr.Slider(minimum=5, maximum=20, value=10, step=1, label='Số lượng kết quả'),
],
outputs=gr.Textbox(lines=20, label='Kết quả'),
title='Vietnamese Legal Document Retrieval System',
description='🔍 Nhập câu hỏi pháp lý của bạn bằng tiếng Việt để nhận các đoạn văn bản pháp luật liên quan.',
examples=[
['Tội xúc phạm danh dự?'],
['Quyền lợi của người lao động?'],
['Thủ tục đăng ký kết hôn?'],
],
flagging_mode='never'
)
if __name__ == '__main__':
demo.launch(share=True) |