Spaces:
Running
Running
File size: 2,313 Bytes
57a0cde | 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 | import os
import torch
from sentence_transformers import SentenceTransformer
# Resilient LangChain Imports
try:
from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever
from langchain_core.documents import Document
except ImportError:
# Fallback for older versions or missing community package
from langchain.retrievers import BM25Retriever, EnsembleRetriever
from langchain.schema import Document
class AdvancedRAG:
"""Advanced RAG system with Hybrid Search (BM25 + Vector) and Reranking."""
def __init__(self, vector_db_collection, embedder_model='all-MiniLM-L6-v2'):
self.collection = vector_db_collection
self.embedder = SentenceTransformer(embedder_model, device='cuda' if torch.cuda.is_available() else 'cpu')
self.bm25_retriever = None
def initialize_bm25(self, texts):
"""Initialize BM25 with a corpus of texts."""
if not texts:
return
documents = [Document(page_content=t) for t in texts]
self.bm25_retriever = BM25Retriever.from_documents(documents)
self.bm25_retriever.k = 5
def hybrid_search(self, query, top_k=5):
"""Perform hybrid search combining vector similarity and keyword matching."""
# 1. Vector Search
query_embedding = self.embedder.encode(query).tolist()
vector_results = self.collection.query(query_embeddings=[query_embedding], n_results=top_k)
vector_docs = vector_results['documents'][0] if vector_results['documents'] else []
# 2. BM25 Search (if initialized)
bm25_docs = []
if self.bm25_retriever:
try:
results = self.bm25_retriever.get_relevant_documents(query)
bm25_docs = [doc.page_content for doc in results]
except: pass
# 3. Ensemble (Simple deduplication and ranking)
combined = list(dict.fromkeys(vector_docs + bm25_docs))
return combined[:top_k]
def add_intelligence(self, new_text):
"""Add new information to the brain."""
embedding = self.embedder.encode(new_text).tolist()
self.collection.add(
documents=[new_text],
embeddings=[embedding],
ids=[str(hash(new_text))]
)
|