vector_rag.py

from sentence_transformers import SentenceTransformer
import faiss
import numpy as np

model = SentenceTransformer("all-MiniLM-L6-v2")

def build_vector(text):
    chunks = [text[i:i+500] for i in range(0, len(text), 500)]
    emb = model.encode(chunks)

    index = faiss.IndexFlatL2(len(emb[0]))
    index.add(np.array(emb))

    return index, chunks


def retrieve(query, index, chunks):
    q = model.encode([query])
    _, I = index.search(q, 3)

    return "\n".join([chunks[i] for i in I[0]])