Update main.py

main.py

@@ -8,9 +8,6 @@ import pandas  # Needed for operating on dataset
 import sentence_transformers  # Needed for query embedding
 import faiss  # Needed for fast similarity search
 
-# Load the model for later use in embeddings
-model = sentence_transformers.SentenceTransformer("allenai-specter")
-
 # Load the dataset and convert to pandas
 full_data = datasets.load_dataset("ccm/publications")["train"].to_pandas()
 
@@ -28,12 +25,18 @@ data.reset_index(inplace=True)
 # Create a FAISS index for fast similarity search
 index = faiss.IndexFlatL2(len(data["embedding"][0]))
 index.metric_type = faiss.METRIC_INNER_PRODUCT
-index.add(faiss.normalize_L2(numpy.stack(data["embedding"].tolist(), axis=0)))
+vectors = numpy.stack(data["embedding"].tolist(), axis=0)
+faiss.normalize_L2(vectors)
+index.add(vectors)
+
+# Load the model for later use in embeddings
+model = sentence_transformers.SentenceTransformer("allenai-specter")
 
 
 # Define the search function
 def search(query: str, k: int):
-    query = numpy.expand_dims(faiss.normalize_L2(model.encode(query)), axis=0)
+    query = numpy.expand_dims(model.encode(query), axis=0)
+    faiss.normalize_L2(query)
     D, I = index.search(query, k)
     top_five = data.loc[I[0]]
     search_results = ""