RAG-Assistant1 / embedder.py
Rixhot's picture
Upload 5 files
afa96cc verified
Raw
History Blame Contribute Delete
1.45 kB
from sentence_transformers import SentenceTransformer
# Load the model once when this file is first imported
# This happens only one time — not every time we call the functions
model = SentenceTransformer("all-MiniLM-L6-v2")
def embed_chunks(chunks):
"""
Converts all chunk texts into vectors.
Input : chunks → list of chunk dicts from chunk_text()
Output : same list but each chunk now has an "embedding" field added
"""
print("Embedding chunks... this may take a minute.")
# Extract just the text from each chunk into a simple list
texts = [chunk["text"] for chunk in chunks]
# Convert all texts to vectors in one batch (faster than one by one)
# show_progress_bar=True shows a progress bar so we know it is working
embeddings = model.encode(texts, show_progress_bar=True)
# Attach each embedding back to its chunk dict
for i, chunk in enumerate(chunks):
chunk["embedding"] = embeddings[i]
print(f"Done! Each chunk is now a vector of {len(embeddings[0])} numbers.")
return chunks
def embed_query(query):
"""
Converts a user question into a vector.
Input : query → a plain string question from the user
Output : a single vector (list of 384 numbers)
We use this at search time to find which chunks are most
similar to what the user is asking.
"""
return model.encode([query])[0] # encode returns a list, [0] gets first item