Initial rework to handle new index format

.gitignore

@@ -0,0 +1 @@
+.venv

README.md

@@ -4,7 +4,7 @@ emoji: 📝
 colorFrom: blue
 colorTo: gray
 sdk: gradio
-sdk_version: 3.29.0
+sdk_version: 5.6.0
 app_file: app.py
 pinned: false
 license: cc0-1.0

app.py

@@ -5,19 +5,36 @@ import requests
 from sentence_transformers import SentenceTransformer
 import faiss
 import gradio as gr
+from datasets import Dataset
+from typing import TypedDict
+import json
 
 from markdown_it import MarkdownIt # used for overriding default markdown renderer
 
+SEARCH_TIME_S = 1  # TODO: optimize
+
+
+class IndexParameters(TypedDict):
+    recall: float  # in this case 10-recall@10
+    exec_time: float  # seconds (raw faiss measure is in milliseconds)
+    param_string: str  # pass directly to faiss index
+
+
 model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
 
-works_ids_path = 'openalex_ids.txt'
-with open(works_ids_path) as f:
-    idxs = f.read().splitlines()
-index = faiss.read_index('index.faiss')
+index: Dataset = Dataset.from_parquet("index/ids.parquet")  # type: ignore
+index.load_faiss_index("embeddings", "index/index.faiss", None)
+faiss_index: faiss.Index = index.get_index("embeddings").faiss_index  # type: ignore
+
+with open("index/params.json", "r") as f:
+    params: list[IndexParameters] = json.load(f)
+params = [p for p in params if p["exec_time"] < SEARCH_TIME_S]
+param = max(params, key=(lambda p: p["recall"]))
+param_string = param["param_string"]
 
 ps = faiss.ParameterSpace()
-ps.initialize(index)
-ps.set_index_parameters(index, 'nprobe=16,ht=512')
+ps.initialize(faiss_index)
+ps.set_index_parameters(faiss_index, param_string)
 
 
 def _recover_abstract(inverted_index):
@@ -33,16 +50,13 @@ def _recover_abstract(inverted_index):
     return abstract
 
 def search(query):
-    global model, index, idxs
+    global model, index
 
+    # TODO: pass in param string directly?
     query_embedding = model.encode(query)
-    query_embedding = query_embedding.reshape(1, -1)
-    distances, faiss_ids = index.search(query_embedding, 20)
-
-    distances = distances[0]
-    faiss_ids = faiss_ids[0]
+    distances, faiss_ids = index.search("embeddings", query_embedding, 20)
 
-    openalex_ids = [idxs[faiss_id] for faiss_id in faiss_ids]
+    openalex_ids = index[faiss_ids]["idxs"]
     search_filter = f'openalex_id:{"|".join(openalex_ids)}'
     search_select = 'id,title,abstract_inverted_index,authorships,primary_location,publication_year,cited_by_count,doi'
 
@@ -148,7 +162,7 @@ with gr.Blocks() as demo:
     response_var = gr.State()
     query = gr.Textbox(lines=1, placeholder='Enter your query here', show_label=False)
     btn = gr.Button('Search')
-    with gr.Box():
+    with gr.Group():
         results = gr.Markdown()
 
     md = MarkdownIt('js-default', {'linkify': True, 'typographer': True}) # don't render html or latex!

index.faiss

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

index.ivfdata

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

openalex_ids.txt

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

requirements.txt

@@ -1,2 +1,3 @@
 sentence-transformers
 faiss-cpu
+datasets