File size: 934 Bytes
71de22d e8954be |
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 |
import pickle
import faiss
import streamlit as st
from sentence_transformers import SentenceTransformer
import pandas as pd
def search(query, model, embeddings, filtered_df, top_x=30):
filtered_df_indecies_list = filtered_df.index
filtered_embeddings = embeddings[filtered_df_indecies_list]
# Load or create FAISS index
dimension = filtered_embeddings.shape[1]
faiss_index = faiss.IndexFlatL2(dimension)
faiss_index.add(filtered_embeddings)
# Convert query to embedding
query_embedding = model.encode([query])[0].reshape(1, -1)
# Perform search
D, I = faiss_index.search(query_embedding, k=top_x) # Search for top x similar items
# Extract the sentences corresponding to the top indices
#print(filtered_df.columns())
top_indecies = [i for i in I[0]]
return filtered_df.iloc[top_indecies] |