Chatbot / knowledge_base.py
Ralqasimi's picture
Update knowledge_base.py
e663f87 verified
raw
history blame
1.2 kB
# Create FAISS index
def create_faiss_index(texts):
"""
Create a FAISS index from the provided list of texts.
"""
import faiss
from sentence_transformers import SentenceTransformer
# Load pre-trained SentenceTransformer model
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
embeddings = model.encode(texts)
# Create the FAISS index
dimension = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(embeddings)
return index, texts
# Search the FAISS index
def search_faiss(faiss_index, stored_texts, query, top_k=3):
"""
Search the FAISS index for the most relevant texts based on the query.
"""
from sentence_transformers import SentenceTransformer
# Load the same model used for indexing
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
# Encode the query into an embedding
query_embedding = model.encode([query])
# Search the FAISS index
distances, indices = faiss_index.search(query_embedding, top_k)
# Retrieve the corresponding texts
results = [stored_texts[i] for i in indices[0] if i < len(stored_texts)]
return results