| |
| |
| |
| from sentence_transformers import SentenceTransformer |
| import chromadb |
| from pathlib import Path |
|
|
| base_dir = Path(__file__).resolve().parent.parent |
| embed_model = SentenceTransformer("multi-qa-mpnet-base-dot-v1") |
| db_path = base_dir / "chroma" / "chroma_db" |
| client = chromadb.PersistentClient(path= str(db_path)) |
|
|
| print(client.list_collections()) |
|
|
| def retrieve_from_vector_db(query: str, collection_name: str = "bank_faq", k: int = 3) -> list[dict]: |
| """ |
| Embeds the query using the same model used during ingestion, |
| queries ChromaDB, and returns chunks with their similarity scores. |
| """ |
|
|
| collection = client.get_collection(name=collection_name) |
|
|
| |
| query_embedding = embed_model.encode(query).tolist() |
|
|
| results = collection.query( |
| query_embeddings=[query_embedding], |
| n_results=k, |
| include=["documents", "metadatas", "distances"] |
| ) |
|
|
| |
| chunks = [] |
| for doc, metadata, distance in zip( |
| results["documents"][0], |
| results["metadatas"][0], |
| results["distances"][0] |
| ): |
| chunks.append({ |
| "text": doc, |
| "metadata": metadata, |
| "similarity_score": round(1 - distance, 4) |
| }) |
|
|
| return chunks |
|
|
| results = retrieve_from_vector_db(query="What are the requirements to open a bank account?", k=3) |
|
|
| |
| |
| |
| |
| |
| |
| |
|
|
|
|
