| import faiss | |
| import gradio as gr | |
| import numpy as np | |
| import pandas as pd | |
| import torch.nn.functional as F | |
| from sentence_transformers import SentenceTransformer | |
| DIM = 768 | |
| model = SentenceTransformer("nomic-ai/nomic-embed-text-v1.5", trust_remote_code=True) | |
| print("Model loaded successfully") | |
| papers_df = pd.read_csv("data/cvpr2024_papers_with_details.csv", index_col=None, on_bad_lines='skip') | |
| papers_df = papers_df[~papers_df["summary"].isna() & ~papers_df["pdf_path"].isna()] | |
| print("Data loaded successfully") | |
| with open('data/embeddings.npy', 'rb') as f: | |
| embeddings = np.load(f) | |
| index = faiss.IndexFlatL2(DIM) | |
| index.add(embeddings) | |
| print("Index loaded successfully") | |
| def encode_query(query): | |
| query_embeddings = model.encode([query], convert_to_tensor=True) | |
| query_embeddings = F.layer_norm(query_embeddings, normalized_shape=(query_embeddings.shape[1],)) | |
| query_embeddings = query_embeddings[:, :DIM] | |
| query_embeddings = F.normalize(query_embeddings, p=2, dim=1) | |
| return query_embeddings | |
| def search_nearest_papers(query, k=5): | |
| query_embeddings = encode_query(query) | |
| D, I = index.search(query_embeddings, k) | |
| return papers_df.iloc[I[0]][["Title", "arXiv_link"]] | |
| demo = gr.Interface( | |
| search_nearest_papers, | |
| [ | |
| "text", | |
| gr.Slider(1, 10, value=5), | |
| ], | |
| gr.Dataframe( | |
| headers=["Title", "PDF"], | |
| ), | |
| title="CVPR 2024 Paper Search", | |
| description="Semantic search over CPVR 2024 paper summary. This app was made using the data available on https://github.com/harpreetsahota204/CVPR-2024-Papers.", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |