Update main.py

main.py

@@ -25,24 +25,27 @@ data.reset_index(inplace=True)
 
 # Create a FAISS index for fast similarity search
 indices = []
-metrics = [faiss.METRIC_INNER_PRODUCT]
+metrics = [faiss.METRIC_INNER_PRODUCT ,faiss.METRIC_L2]
+normalization = [True, False]
 vectors = numpy.stack(data["embedding"].tolist(), axis=0)
 for  metric in metrics:
-    index = faiss.IndexFlatL2(len(data["embedding"][0]))
-    index.metric_type = metric
-    faiss.normalize_L2(vectors)
-    index.train(vectors)
-    index.add(vectors)
-    indices.append(index)
+    for normal in normalization
+        index = faiss.IndexFlatL2(len(data["embedding"][0]))
+        index.metric_type = metric
+        if normal:
+            faiss.normalize_L2(vectors)
+        index.train(vectors)
+        index.add(vectors)
+        indices.append(index)
 
 # Load the model for later use in embeddings
 model = sentence_transformers.SentenceTransformer("allenai-specter")
 
 # Define the search function
-def search(query: str, k: int):
+def search(query: str, k: int, n: int):
     query = numpy.expand_dims(model.encode(query), axis=0)
     faiss.normalize_L2(query)
-    D, I = indices[0].search(query, k)
+    D, I = indices[n].search(query, k)
     top_five = data.loc[I[0]]
     search_results = ""
 
@@ -88,8 +91,10 @@ with gradio.Blocks() as demo:
         )
         with gradio.Accordion("Settings", open=False):
             k = gradio.Number(10.0, label="Number of results", precision=0)
+            k = gradio.Radio([True, False], label="Number of results", precision=0)
     results = gradio.Markdown()
-    query.change(fn=search, inputs=[query, k], outputs=results)
-    k.change(fn=search, inputs=[query, k], outputs=results)
+    query.change(fn=search, inputs=[query, k, n], outputs=results)
+    k.change(fn=search, inputs=[query, k, n], outputs=results)
+    n.change(fn=search, inputs=[query, k, n], outputs=results)
 
 demo.launch(debug=True)