| import faiss |
| import numpy as np |
|
|
| from embeddings import get_embeddings |
|
|
|
|
| class RAG: |
|
|
| def __init__(self): |
| self.index = None |
| self.chunks = [] |
|
|
| def chunk_text(self, text, chunk_size=500): |
|
|
| words = text.split() |
|
|
| chunks = [] |
|
|
| for i in range(0, len(words), chunk_size): |
| chunks.append(" ".join(words[i:i + chunk_size])) |
|
|
| return chunks |
|
|
| def create_index(self, text): |
|
|
| self.chunks = self.chunk_text(text) |
| if not self.chunks: |
| return |
|
|
| embeddings = get_embeddings(self.chunks) |
|
|
| embeddings = np.atleast_2d(np.array(embeddings)).astype("float32") |
|
|
| dimension = embeddings.shape[1] |
|
|
| self.index = faiss.IndexFlatL2(dimension) |
|
|
| self.index.add(embeddings) |
|
|
| def search(self, query, top_k=3): |
| if self.index is None or not self.chunks: |
| return [] |
|
|
| query_embedding = get_embeddings([query]) |
| query_embedding = np.atleast_2d(np.array(query_embedding)).astype("float32") |
|
|
| top_k = min(top_k, len(self.chunks)) |
| distances, indices = self.index.search(query_embedding, top_k) |
|
|
| results = [] |
|
|
| for score, idx in zip(distances[0], indices[0]): |
|
|
| if idx == -1: |
| continue |
|
|
| results.append( |
| { |
| "chunk": self.chunks[idx], |
| "score": float(score), |
| "chunk_id": int(idx) |
| } |
| ) |
|
|
| return results |
|
|
