File size: 7,040 Bytes
ba1724a 7285400 ba1724a 7285400 da3c141 ba1724a 7285400 ba1724a 7285400 ba1724a 7285400 ba1724a 7285400 ba1724a 7285400 ba1724a 7285400 |
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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 |
import re
import gradio as gr
from scipy.sparse import load_npz
import torch
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.preprocessing import normalize
from transformers import BertTokenizer, BertModel
import numpy as np
from datasets import load_dataset
from gensim.models import KeyedVectors
class ArxivSearch:
def __init__(self, dataset, embedding="tfidf"):
self.dataset = dataset
self.embedding = embedding
self.documents = []
self.titles = []
self.raw_texts = []
self.arxiv_ids = []
self.embedding_dropdown = gr.Dropdown(
choices=["tfidf", "word2vec", "bert"],
value="tfidf",
label="Model"
)
self.iface = gr.Interface(
fn=self.search_function,
inputs=[
gr.Textbox(lines=1, placeholder="Enter your search query"),
self.embedding_dropdown
],
outputs=gr.Markdown(),
title="arXiv Search Engine",
description="Search arXiv papers by keyword and embedding model.",
)
self.load_data(dataset)
self.load_model(embedding)
self.iface.launch()
# # --- Load data and embeddings ---
# with open("feature_names.txt", "r") as f:
# feature_names = [line.strip() for line in f]
# tfidf_matrix = load_npz("tfidf_matrix_train.npz")
# Load dataset and initialize search engine
def load_data(self, dataset):
train_data = dataset["train"]
for item in train_data.select(range(len(train_data))):
text = item["text"]
if not text or len(text.strip()) < 10:
continue
lines = text.splitlines()
title_lines = []
found_arxiv = False
arxiv_id = None
for line in lines:
line_strip = line.strip()
if not found_arxiv and line_strip.lower().startswith("arxiv:"):
found_arxiv = True
match = re.search(r'arxiv:\d{4}\.\d{4,5}v\d', line_strip, flags=re.IGNORECASE)
if match:
arxiv_id = match.group(0).lower()
elif not found_arxiv:
title_lines.append(line_strip)
else:
if line_strip.lower().startswith("abstract"):
break
title = " ".join(title_lines).strip()
self.raw_texts.append(text.strip())
self.titles.append(title)
self.documents.append(text.strip())
self.arxiv_ids.append(arxiv_id)
def keyword_match_ranking(self, query, top_n=5):
query_terms = query.lower().split()
query_indices = [i for i, term in enumerate(self.feature_names) if term in query_terms]
if not query_indices:
return []
scores = []
for doc_idx in range(self.tfidf_matrix.shape[0]):
doc_vector = self.tfidf_matrix[doc_idx]
doc_score = sum(doc_vector[0, i] for i in query_indices)
if doc_score > 0:
scores.append((doc_idx, doc_score))
scores.sort(key=lambda x: x[1], reverse=True)
return scores[:top_n]
def word2vec_search(self, query, top_n=5):
tokens = [word for word in query.split() if word in self.wv_model.key_to_index]
if not tokens:
return []
vectors = np.array([self.wv_model[word] for word in tokens])
query_vec = normalize(np.mean(vectors, axis=0).reshape(1, -1))
sims = cosine_similarity(query_vec, self.word2vec_embeddings).flatten()
top_indices = sims.argsort()[::-1][:top_n]
return [(i, sims[i]) for i in top_indices]
def bert_search(self, query, top_n=5):
with torch.no_grad():
inputs = self.tokenizer(query, return_tensors="pt", truncation=True, padding=True)
outputs = self.model(**inputs)
query_vec = normalize(outputs.last_hidden_state[:, 0, :].numpy())
sims = cosine_similarity(query_vec, self.bert_embeddings).flatten()
top_indices = sims.argsort()[::-1][:top_n]
return [(i, sims[i]) for i in top_indices]
def load_model(self, embedding):
if embedding == "tfidf":
self.tfidf_matrix = load_npz("TF-IDF embeddings/tfidf_matrix_train.npz")
with open("TF-IDF embeddings/feature_names.txt", "r") as f:
self.feature_names = [line.strip() for line in f.readlines()]
elif embedding == "word2vec":
# Use trimmed model here
self.word2vec_embeddings = normalize(np.load("Word2Vec embeddings/word2vec_embedding.npz")["word2vec_embedding"])
self.wv_model = KeyedVectors.load("models/word2vec-trimmed.model")
elif embedding == "bert":
self.bert_embeddings = normalize(np.load("BERT embeddings/bert_embedding.npz")["bert_embedding"])
self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
self.model = BertModel.from_pretrained('bert-base-uncased')
self.model.eval()
else:
raise ValueError(f"Unsupported embedding type: {embedding}")
def on_model_change(self, change):
new_model = change["new"]
self.embedding = new_model
self.load_model(new_model)
def snippet_before_abstract(self, text):
pattern = re.compile(r'a\s*b\s*s\s*t\s*r\s*a\s*c\s*t|i\s*n\s*t\s*r\s*o\s*d\s*u\s*c\s*t\s*i\s*o\s*n', re.IGNORECASE)
match = pattern.search(text)
if match:
return text[:match.start()].strip()
else:
return text[:100].strip()
def search_function(self, query, embedding):
# Load or switch embedding model here if needed
if embedding == "tfidf":
results = self.keyword_match_ranking(query)
elif embedding == "word2vec":
results = self.word2vec_search(query)
elif embedding == "bert":
results = self.bert_search(query)
else:
return "No results found."
if not results:
return "No results found."
output = ""
display_rank = 1
for idx, score in results:
if not self.arxiv_ids[idx]:
continue
link = f"https://arxiv.org/abs/{self.arxiv_ids[idx].replace('arxiv:', '')}"
snippet = self.snippet_before_abstract(self.documents[idx]).replace('\n', '<br>')
output += f"### Document {display_rank}\n"
output += f"[arXiv Link]({link})\n\n"
output += f"<pre>{snippet}</pre>\n\n---\n"
display_rank += 1
return output
if __name__ == "__main__":
dataset = load_dataset("ccdv/arxiv-classification", "no_ref") # replace with your dataset
search_engine = ArxivSearch(dataset, embedding="tfidf") # Initialize with tfidf or any other embedding
search_engine.iface.launch() |