File size: 1,507 Bytes
959c484 | 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 | 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
|