| 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=20): |
|
|
| filtered_df_indecies_list = filtered_df.index |
| filtered_embeddings = embeddings[filtered_df_indecies_list] |
| |
| |
| dimension = filtered_embeddings.shape[1] |
| faiss_index = faiss.IndexFlatL2(dimension) |
| faiss_index.add(filtered_embeddings) |
|
|
| |
| query_embedding = model.encode([query])[0].reshape(1, -1) |
| |
| |
| D, I = faiss_index.search(query_embedding, k=top_x) |
| |
| |
| |
| top_indecies = [i for i in I[0]] |
|
|
| return filtered_df.iloc[top_indecies] |