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
import pandas as pd
from datasets import load_dataset
from gensim.models import KeyedVectors
import plotly.graph_objects as go
from sklearn.decomposition import PCA
from transformers import AutoTokenizer, AutoModel
from sentence_transformers import CrossEncoder
from sentence_transformers import SentenceTransformer
class ArxivSearch:
def __init__(self, dataset, embedding="sbert"):
self.dataset = dataset
self.embedding = embedding
self.query = None
self.documents = []
self.titles = []
self.raw_texts = []
self.arxiv_ids = []
self.last_results = []
self.query_encoding = None
# model selection
self.embedding_dropdown = gr.Dropdown(
choices=["tfidf", "word2vec", "bert", "sbert", "clustered sbert"],
value="sbert",
label="Model"
)
self.plot_button = gr.Button("Show 3D Plot")
# Gradio blocks for UI elements
with gr.Blocks() as self.iface:
gr.Markdown("# arXiv Search Engine")
gr.Markdown("Search arXiv papers by keyword and embedding model.")
self.plot_output = gr.Plot()
with gr.Row():
self.query_box = gr.Textbox(lines=1, placeholder="Enter your search query", label="Query")
self.embedding_dropdown.render()
self.plot_button.render()
with gr.Column():
self.search_button = gr.Button("Search")
self.output_md = gr.Markdown()
self.query_box.submit(
self.search_function,
inputs=[self.query_box, self.embedding_dropdown],
outputs=self.output_md
)
# self.embedding_dropdown.change(
# self.model_switch,
# inputs=[self.embedding_dropdown],
# outputs=self.output_md
# )
self.embedding_dropdown.change(
self.search_function,
inputs=[self.query_box, self.embedding_dropdown],
outputs=self.output_md
)
self.plot_button.click(
self.plot_3d_embeddings,
inputs=[],
outputs=self.plot_output
)
self.search_button.click(
self.search_function,
inputs=[self.query_box, self.embedding_dropdown],
outputs=self.output_md
)
self.load_data(dataset)
# self.load_model(embedding)
self.load_model('tfidf')
self.load_model('word2vec')
self.load_model('bert')
# self.load_model('scibert')
# self.load_model('sbert')
self.load_model('clustered sbert')
self.iface.launch()
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 plot_dense(self, embedding, pca, results_indices):
all_indices = list(set(results_indices) | set(range(min(5000, embedding.shape[0]))))
all_data = embedding[all_indices]
pca.fit(all_data)
reduced_data = pca.transform(embedding[:5000])
reduced_results_points = pca.transform(embedding[results_indices]) if len(results_indices) > 0 else np.empty((0, 3))
query_point = pca.transform(self.query_encoding) if self.query_encoding is not None and self.query_encoding.shape[0] > 0 else np.empty((0, 3))
return reduced_data, reduced_results_points, query_point
def plot_3d_embeddings(self):
# Example: plot random points, replace with your embeddings
pca = PCA(n_components=3)
results_indices = [i[0] for i in self.last_results]
if self.embedding == "tfidf":
all_indices = list(set(results_indices) | set(range(min(5000, self.tfidf_matrix.shape[0]))))
all_data = self.tfidf_matrix[all_indices].toarray()
pca.fit(all_data)
reduced_data = pca.transform(self.tfidf_matrix[:5000].toarray())
reduced_results_points = pca.transform(self.tfidf_matrix[results_indices].toarray()) if len(results_indices) > 0 else np.empty((0, 3))
elif self.embedding == "word2vec":
reduced_data, reduced_results_points, query_point = self.plot_dense(self.word2vec_embeddings, pca, results_indices)
elif self.embedding == "bert":
reduced_data, reduced_results_points, query_point = self.plot_dense(self.bert_embeddings, pca, results_indices)
elif self.embedding == "sbert" or self.embedding == "clustered sbert":
reduced_data, reduced_results_points, query_point = self.plot_dense(self.sbert_embedding, pca, results_indices)
if self.embedding == "clustered sbert":
cluster_colors = ["#00b7ff" if i in np.where(self.clusters == self.top_cluster_index)[0] else "#ffffff" for i in range(len(self.documents))]
# elif self.embedding == "scibert":
# reduced_data, reduced_results_points, query_point = self.plot_dense(self.scibert_embeddings, pca, results_indices)
else:
raise ValueError(f"Unsupported embedding type: {self.embedding}")
results_scores = [i[1] for i in self.last_results]
traces = []
trace = go.Scatter3d(
x=reduced_data[:, 0],
y=reduced_data[:, 1],
z=reduced_data[:, 2],
mode='markers',
marker=dict(size=3.5,
color="#ffffff" if self.embedding != "clustered sbert" else cluster_colors,
opacity=0.2),
name='All Documents',
text=[f"
: {self.arxiv_ids[i] if self.arxiv_ids[i] else self.documents[i].split()[:10]}" for i in range(len(self.documents))],
hoverinfo='text'
)
traces.append(trace)
layout = go.Layout(
margin=dict(l=0, r=0, b=0, t=0),
scene=dict(
xaxis_title='PCA 1',
yaxis_title='PCA 2',
zaxis_title='PCA 3',
xaxis=dict(backgroundcolor='black', color='white', gridcolor='gray', zerolinecolor='gray'),
yaxis=dict(backgroundcolor='black', color='white', gridcolor='gray', zerolinecolor='gray'),
zaxis=dict(backgroundcolor='black', color='white', gridcolor='gray', zerolinecolor='gray'),
),
paper_bgcolor='black', # Outside the plotting area
plot_bgcolor='black', # Plotting area
font=dict(color='white'), # Axis and legend text
legend=dict(
bgcolor='rgba(0,0,0,0)', # Transparent legend background
bordercolor='rgba(0,0,0,0)', # No border
x=0.01, # Place legend inside plot area (adjust as needed)
y=0.99,
xanchor='left',
yanchor='top'
)
)
if len(reduced_results_points) > 0:
custom_colorscale = [
[0.0, "#00ffea"], # Start color (e.g., bright cyan)
[1.0, "#ffea00"], # End color (e.g., bright yellow)
]
results_trace = go.Scatter3d(
x=reduced_results_points[:, 0],
y=reduced_results_points[:, 1],
z=reduced_results_points[:, 2],
mode='markers',
marker=dict(size=4.25,
color=results_scores,
colorscale=custom_colorscale,
opacity=0.99,
colorbar=dict(
title="Score",
bgcolor='rgba(0,0,0,0)', # <-- Transparent background for colorbar
bordercolor='rgba(0,0,0,0)' # No border
)
),
name='Results',
text=[f"
{self.documents[i][:100]}" for i in results_indices],
hoverinfo='text'
)
traces.append(results_trace)
if not self.embedding == "tfidf" and self.query_encoding is not None and self.query_encoding.shape[0] > 0:
query_trace = go.Scatter3d(
x=query_point[:, 0],
y=query_point[:, 1],
z=query_point[:, 2],
mode='markers',
marker=dict(size=5, color='red', opacity=0.8),
name='Query',
text=[f"
Query: {self.query}"],
hoverinfo='text'
)
traces.append(query_trace)
fig = go.Figure(data=traces, layout=layout)
return fig
def keyword_match_ranking(self, query, top_n=10):
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=10):
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 = np.mean(vectors, axis=0).reshape(1, -1)
self.query_encoding = query_vec
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=10):
with torch.no_grad():
inputs = self.tokenizer((query+' ')*2, return_tensors="pt", truncation=True, max_length=512, padding='max_length')
outputs = self.model(**inputs)
query_vec = outputs.last_hidden_state[:, 0, :].numpy()
self.query_encoding = query_vec
sims = cosine_similarity(query_vec, self.bert_embeddings).flatten()
top_indices = sims.argsort()[::-1][:top_n]
print(f"sim, top_indices: {sims}, {top_indices}")
return [(i, sims[i]) for i in top_indices]
# def scibert_search(self, query, top_n=10):
# with torch.no_grad():
# inputs = self.sci_tokenizer(query, return_tensors="pt", truncation=True, padding=True, max_length=512)
# outputs = self.sci_model(**inputs)
# query_vec = outputs.last_hidden_state[:, 0, :].numpy()
# self.query_encoding = query_vec
# sims = cosine_similarity(query_vec, self.scibert_embeddings).flatten()
# top_indices = sims.argsort()[::-1][:top_n]
# print(f"sim, top_indices: {sims}, {top_indices}")
# return [(i, sims[i]) for i in top_indices]
def sbert_search(self, query, top_n=10):
query_vec = self.sbert_model.encode([query])
self.query_encoding = query_vec
cos_scores = cosine_similarity(query_vec, self.sbert_embedding)[0]
top_k_indices = np.argsort(cos_scores)[-50:][::-1]
candidates = [dataset['train'][int(i)]['text'] for i in top_k_indices]
scores = self.cross_encoder.predict([(query, doc) for doc in candidates])
final_scores = 0.7 * scores + 0.3 * cos_scores[top_k_indices]
top_indices = top_k_indices[final_scores.argsort()[::-1][:top_n]]
print(f"sim, top_indices: {final_scores}, {top_indices}")
return [(top_k_indices[i], final_scores[i]) for i in final_scores.argsort()[::-1][:top_n]]
def clustered_sbert_search(self, query, top_n=10):
query_vec = self.sbert_model.encode([query])
self.query_encoding = query_vec # Store the query encoding for plotting
cos_cluster_scores = cosine_similarity(query_vec, self.cluster_centers)[0] # Get cosine similarity with cluster centers
self.top_cluster_index = np.argmax(cos_cluster_scores) # Get the index of the top cluster
cos_scores = cosine_similarity(query_vec, self.clustered_embeddings[self.top_cluster_index])[0] # Get cosine similarity within the top cluster
top_k_indices = np.argsort(cos_scores)[-50:][::-1] # Get top 50 indices within the top cluster (cluster internal indices)
top_full_dataset_indices = np.where(self.clusters == self.top_cluster_index)[0][top_k_indices] # Get the 50 indices that correspond to the full dataset
candidates = [self.dataset['train'][int(i)]['text'] for i in top_full_dataset_indices] # Get the 50 candidate documents
scores = self.cross_encoder.predict([(query, doc) for doc in candidates]) # Get the 50 cross-encoder scores for the candidates
final_scores = 0.7 * scores + 0.3 * cos_scores[top_k_indices] # Combine the 50 cross-encoder scores with the cosine similarity scores
top_indices = top_k_indices[final_scores.argsort()[::-1][:top_n]] # Get the top N cluster internal indices based on the final scores
top_indices_full = np.where(self.clusters == self.top_cluster_index)[0][top_indices] # Get the top N full dataset indices based on the final scores
print(f"sim, top_indices: {final_scores}, {top_indices}")
return [(i, final_scores[j]) for j, i in enumerate(top_indices_full)]
def model_switch(self, embedding, progress=gr.Progress()):
if self.embedding != embedding:
old_embedding = self.embedding
print(f"Switching model to {embedding}")
self.load_model(embedding)
print(f"Loaded {embedding} model")
self.embedding = embedding
if old_embedding == "tfidf":
del self.tfidf_matrix
del self.feature_names
if old_embedding == "word2vec":
del self.word2vec_embeddings
del self.wv_model
if old_embedding == "bert":
del self.bert_embeddings
del self.tokenizer
del self.model
if old_embedding == "scibert":
del self.scibert_embeddings
del self.sci_tokenizer
del self.sci_model
if old_embedding == "sbert":
del self.sbert_model
del self.sbert_embedding
del self.cross_encoder
print(f"old embedding removed")
if hasattr(self, "query") and self.query:
return self.search_function(self.query, self.embedding)
else:
return "" # Or a message like "Model switched. Please enter a query."
return gr.update() # No change if embedding is the same
def load_model(self, embedding):
self.embedding = embedding
if self.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 self.embedding == "word2vec":
# Use trimmed model here
self.word2vec_embeddings = np.load("Word2Vec embeddings/word2vec_embedding.npz")["word2vec_embedding"]
self.wv_model = KeyedVectors.load("models/word2vec-trimmed.model")
elif self.embedding == "bert":
self.bert_embeddings = 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()
# elif self.embedding == "scibert":
# self.scibert_embeddings = np.load("SciBERT_embeddings/scibert_embedding.npz")["bert_embedding"]
# self.sci_tokenizer = AutoTokenizer.from_pretrained('allenai/scibert_scivocab_uncased')
# self.sci_model = AutoModel.from_pretrained('allenai/scibert_scivocab_uncased')
# self.sci_model.eval()
elif self.embedding == "sbert" or self.embedding == "clustered sbert":
self.sbert_model = SentenceTransformer("all-MiniLM-L6-v2")
self.sbert_embedding = np.load("BERT embeddings/sbert_embedding.npz")["sbert_embedding"]
# self.cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")
self.cross_encoder = CrossEncoder("cross-encoder/ms-marco-TinyBERT-L-2-v2")
if self.embedding == "clustered sbert":
self.clusters = pd.read_csv(f'raf_clusters/cluster_labels_sbert.csv')['cluster_label'].values
self.cluster_centers = pd.read_csv(f'BERT embeddings/sbert_cluster_centers.csv').values
self.clustered_embeddings = [self.sbert_embedding[self.clusters == i] for i in np.unique(self.clusters)]
else:
raise ValueError(f"Unsupported embedding type: {self.embedding}")
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() if match.start() < 1000 else text[:100].strip()
else:
return text[:300].strip()
def set_embedding(self, embedding):
self.embedding = embedding
def search_function(self, query, embedding, progress=gr.Progress()):
self.set_embedding(embedding)
self.query = query
query = query.encode().decode('unicode_escape') # Interpret escape sequences
search_methods = {
"tfidf": self.keyword_match_ranking,
"word2vec": self.word2vec_search,
"bert": self.bert_search,
# "scibert": self.scibert_search, # Uncomment if implemented
"sbert": self.sbert_search,
"clustered sbert": self.clustered_sbert_search,
}
results = search_methods.get(self.embedding, lambda q: [])(query)
if not results:
self.last_results = []
return "No results found."
if results:
self.last_results = results
output = ""
display_rank = 1
for idx, score in results:
if not self.arxiv_ids[idx]:
output += f"### Document {display_rank}\n"
output += f"{self.documents[idx][:200]}\n\n"
else:
link = f"https://arxiv.org/abs/{self.arxiv_ids[idx].replace('arxiv:', '')}"
snippet = self.snippet_before_abstract(self.documents[idx]).replace('\n', '
')
output += f"### Document {display_rank}\n"
output += f"[arXiv Link]({link})\n\n"
output += f"{snippet}\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)
search_engine.iface.launch()