fixed bert not finding documents

SciBERT_embeddings/scibert_embedding.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb6a12b25db606552aba11f26f9bc5c2ac475183b93a21c6abddc743087e3bcd
+size 80887259

app.py

@@ -10,14 +10,15 @@ 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
 
 
 
 class ArxivSearch:
-    def __init__(self, dataset, embedding="tfidf"):
+    def __init__(self, dataset, embedding="bert"):
         self.dataset = dataset
         self.embedding = embedding
+        self.query = None
         self.documents = []
         self.titles = []
         self.raw_texts = []
@@ -25,17 +26,16 @@ class ArxivSearch:
         self.last_results = []
         self.query_encoding = None
 
+        # model selection
         self.embedding_dropdown = gr.Dropdown(
-            choices=["tfidf", "word2vec", "bert"],
-            value="tfidf",
+            choices=["tfidf", "word2vec", "bert", "scibert"],
+            value="bert",
             label="Model"
             )
-        
-
-                # Add a button to show the 3D plot
+    
         self.plot_button = gr.Button("Show 3D Plot")
 
-        # Define the interface using Blocks for more flexibility
+        # 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.")
@@ -64,7 +64,7 @@ class ArxivSearch:
             )
             self.plot_button.click(
                 self.plot_3d_embeddings,
-                inputs=[self.embedding_dropdown],
+                inputs=[],
                 outputs=self.plot_output
             )
             self.search_button.click(
@@ -73,22 +73,11 @@ class ArxivSearch:
                 outputs=self.output_md
             )
 
-        # 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.load_model('tfidf')
         self.load_model('word2vec')
         self.load_model('bert')
+        self.load_model('scibert')
 
         self.iface.launch()
 
@@ -124,19 +113,18 @@ class ArxivSearch:
             self.documents.append(text.strip())
             self.arxiv_ids.append(arxiv_id)
 
-
-    def plot_3d_embeddings(self, embedding):
+    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 embedding == "tfidf":
+        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 embedding == "word2vec":
+        elif self.embedding == "word2vec":
             all_indices = list(set(results_indices) | set(range(min(5000, self.word2vec_embeddings.shape[0]))))
             all_data = self.word2vec_embeddings[all_indices]
             pca.fit(all_data)
@@ -144,16 +132,22 @@ class ArxivSearch:
             reduced_results_points = pca.transform(self.word2vec_embeddings[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))
 
-        elif embedding == "bert":
+        elif self.embedding == "bert":
             all_indices = list(set(results_indices) | set(range(min(5000, self.bert_embeddings.shape[0]))))
             all_data = self.bert_embeddings[all_indices]
             pca.fit(all_data)
             reduced_data = pca.transform(self.bert_embeddings[:5000])
             reduced_results_points = pca.transform(self.bert_embeddings[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))
-
+        elif self.embedding == "scibert":
+            all_indices = list(set(results_indices) | set(range(min(5000, self.scibert_embeddings.shape[0]))))
+            all_data = self.scibert_embeddings[all_indices]
+            pca.fit(all_data)
+            reduced_data = pca.transform(self.scibert_embeddings[:5000])
+            reduced_results_points = pca.transform(self.scibert_embeddings[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))
         else:
-            raise ValueError(f"Unsupported embedding type: {embedding}")
+            raise ValueError(f"Unsupported embedding type: {self.embedding}")
         trace = go.Scatter3d(
             x=reduced_data[:, 0],
             y=reduced_data[:, 1],
@@ -185,7 +179,7 @@ class ArxivSearch:
                 marker=dict(size=3.5, color='orange', opacity=0.75),
                 name='Results'
             )
-            if not "tfidf" and self.query_encoding is not None and self.query_encoding.shape[0] > 0:
+            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],
@@ -231,11 +225,23 @@ class ArxivSearch:
             inputs = self.tokenizer(query, return_tensors="pt", truncation=True, padding=True)
             outputs = self.model(**inputs)
             query_vec = normalize(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]
         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)
+            outputs = self.sci_model(**inputs)
+            query_vec = normalize(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]
+        return [(i, sims[i]) for i in top_indices]
+
     def bert_search_2(self, query, top_n=10):
         with torch.no_grad():
             inputs = self.tokenizer(query, return_tensors="pt", truncation=True, padding=True)
@@ -252,27 +258,28 @@ class ArxivSearch:
         return [(i, sims[i]) for i in top_indices]
 
     def load_model(self, embedding):
-        if embedding == "tfidf":
+        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 embedding == "word2vec":
+        elif self.embedding == "word2vec":
             # Use trimmed model here
-            self.word2vec_embeddings = normalize(np.load("Word2Vec embeddings/word2vec_embedding.npz")["word2vec_embedding"])
+            self.word2vec_embeddings = 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"])
+        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()
         else:
-            raise ValueError(f"Unsupported embedding type: {embedding}")
+            raise ValueError(f"Unsupported embedding type: {self.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)
@@ -284,16 +291,18 @@ class ArxivSearch:
 
 
     def search_function(self, query, embedding):
-        # Preprocess the query
-        query = query.strip().lower()
+        self.embedding = embedding
+        query = query.encode().decode('unicode_escape')  # Interpret escape sequences
 
         # Load or switch embedding model here if needed
-        if embedding == "tfidf":
+        if self.embedding == "tfidf":
             results = self.keyword_match_ranking(query)
-        elif embedding == "word2vec":
+        elif self.embedding == "word2vec":
             results = self.word2vec_search(query)
-        elif embedding == "bert":
+        elif self.embedding == "bert":
             results = self.bert_search(query)
+        elif self.embedding == "scibert":
+            results = self.scibert_search(query)
         else:
             return "No results found."
 
@@ -301,7 +310,6 @@ class ArxivSearch:
             self.last_results = []
             return "No results found."
         
-
         if results:
             self.last_results = results
 
@@ -323,5 +331,5 @@ class ArxivSearch:
 
 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 = ArxivSearch(dataset)
     search_engine.iface.launch()