Spaces:
Sleeping
Sleeping
File size: 3,428 Bytes
bb723d2 | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 | from huggingface_hub import hf_hub_download
from sentence_transformers import SentenceTransformer
import faiss
import json
import pickle
import numpy as np
repoId = "negi2725/dataRag"
repoType = "dataset"
encoder = SentenceTransformer("BAAI/bge-large-en-v1.5")
constitutionIndexPath = hf_hub_download(repo_id=repoId, repo_type=repoType, filename="constitution_bgeLarge.index")
ipcIndexPath = hf_hub_download(repo_id=repoId, repo_type=repoType, filename="ipc_bgeLarge.index")
ipcCaseIndexPath = hf_hub_download(repo_id=repoId, repo_type=repoType, filename="ipc_case_flat.index")
statuteIndexPath = hf_hub_download(repo_id=repoId, repo_type=repoType, filename="statute_index.faiss")
qaIndexPath = hf_hub_download(repo_id=repoId, repo_type=repoType, filename="qa_faiss_index.idx")
caseIndexPath = hf_hub_download(repo_id=repoId, repo_type=repoType, filename="case_faiss.index")
constitutionChunksPath = hf_hub_download(repo_id=repoId, repo_type=repoType, filename="constitution_chunks.json")
ipcChunksPath = hf_hub_download(repo_id=repoId, repo_type=repoType, filename="ipc_chunks.json")
ipcCaseChunksPath = hf_hub_download(repo_id=repoId, repo_type=repoType, filename="ipc_case_chunks.json")
qaChunksPath = hf_hub_download(repo_id=repoId, repo_type=repoType, filename="qa_text_chunks.json")
statuteChunksPath = hf_hub_download(repo_id=repoId, repo_type=repoType, filename="statute_chunks.pkl")
caseChunksPath = hf_hub_download(repo_id=repoId, repo_type=repoType, filename="case_chunks.pkl")
constitutionIndex = faiss.read_index(constitutionIndexPath)
ipcIndex = faiss.read_index(ipcIndexPath)
ipcCaseIndex = faiss.read_index(ipcCaseIndexPath)
statuteIndex = faiss.read_index(statuteIndexPath)
qaIndex = faiss.read_index(qaIndexPath)
caseIndex = faiss.read_index(caseIndexPath)
with open(constitutionChunksPath, "r") as f:
constitutionChunks = json.load(f)
with open(ipcChunksPath, "r") as f:
ipcChunks = json.load(f)
with open(ipcCaseChunksPath, "r") as f:
ipcCaseChunks = json.load(f)
with open(qaChunksPath, "r") as f:
qaChunks = json.load(f)
with open(statuteChunksPath, "rb") as f:
statuteChunks = pickle.load(f)
with open(caseChunksPath, "rb") as f:
caseChunks = pickle.load(f)
def retrieve(text: str, topK: int = 5) -> dict:
queryEmbedding = encoder.encode([text])
queryEmbedding = queryEmbedding.astype("float32")
faiss.normalize_L2(queryEmbedding)
results = {}
distances, indices = constitutionIndex.search(queryEmbedding, topK)
results["constitution"] = [constitutionChunks[idx] for idx in indices[0] if idx < len(constitutionChunks)]
distances, indices = ipcIndex.search(queryEmbedding, topK)
results["ipc"] = [ipcChunks[idx] for idx in indices[0] if idx < len(ipcChunks)]
distances, indices = ipcCaseIndex.search(queryEmbedding, topK)
results["ipcCase"] = [ipcCaseChunks[idx] for idx in indices[0] if idx < len(ipcCaseChunks)]
distances, indices = statuteIndex.search(queryEmbedding, topK)
results["statute"] = [statuteChunks[idx] for idx in indices[0] if idx < len(statuteChunks)]
distances, indices = qaIndex.search(queryEmbedding, topK)
results["qa"] = [qaChunks[idx] for idx in indices[0] if idx < len(qaChunks)]
distances, indices = caseIndex.search(queryEmbedding, topK)
results["case"] = [caseChunks[idx] for idx in indices[0] if idx < len(caseChunks)]
return results
|