Create and clean up function declarations

app.py

@@ -8,9 +8,12 @@ import gradio as gr
 from datasets import Dataset
 from typing import TypedDict
 import json
+from pathlib import Path
 
 from markdown_it import MarkdownIt # used for overriding default markdown renderer
 
+MODEL_NAME = "all-MiniLM-L6-v2"  # TODO: make configurable
+DIR = Path("index")
 SEARCH_TIME_S = 1  # TODO: optimize
 
 
@@ -20,24 +23,29 @@ class IndexParameters(TypedDict):
     param_string: str  # pass directly to faiss index
 
 
-model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
+def get_model(model_name: str, device: str) -> SentenceTransformer:
+    return SentenceTransformer(model_name, device=device)
 
-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"]
+def get_index(dir: Path, search_time_s: float) -> Dataset:
+    index: Dataset = Dataset.from_parquet(str(dir / "ids.parquet"))  # type: ignore
+    index.load_faiss_index("embeddings", dir / "index.faiss", None)
+    faiss_index: faiss.Index = index.get_index("embeddings").faiss_index  # type: ignore
 
-ps = faiss.ParameterSpace()
-ps.initialize(faiss_index)
-ps.set_index_parameters(faiss_index, param_string)
+    with open(dir / "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(faiss_index)
+    ps.set_index_parameters(faiss_index, param_string)
 
-def _recover_abstract(inverted_index):
+    return index
+
+
+def recover_abstract(inverted_index):
     abstract_size = max([max(appearances) for appearances in inverted_index.values()])+1
 
     abstract = [None]*abstract_size
@@ -49,26 +57,12 @@ def _recover_abstract(inverted_index):
     abstract = ' '.join(abstract)
     return abstract
 
-def search(query):
-    global model, index
-
-    # TODO: pass in param string directly?
-    query_embedding = model.encode(query)
-    distances, faiss_ids = index.search("embeddings", query_embedding, 20)
-
-    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'
-
-    neighbors = [(distance, openalex_id) for distance, openalex_id in zip(distances, openalex_ids)]
-    request_str = f'https://api.openalex.org/works?filter={search_filter}&select={search_select}'
-
-    return neighbors, request_str
 
 def execute_request(request_str):
     response = requests.get(request_str).json()
     return response
 
+
 def format_response(neighbors, response):
     response = {doc['id']: doc for doc in response['results']}
     
@@ -92,7 +86,7 @@ def format_response(neighbors, response):
         if abstract_inverted_index is None: # edge case: no abstract
             abstract = 'No abstract'
         else:
-            abstract = _recover_abstract(abstract_inverted_index)
+            abstract = recover_abstract(abstract_inverted_index)
             abstract = abstract.replace('\n', '\\n').replace('\r', '\\r')
 
         # try to get journal name or else set it to None
@@ -146,6 +140,27 @@ def format_response(neighbors, response):
     
     return result_string
 
+
+model = get_model(MODEL_NAME, "cpu")
+index = get_index(DIR, SEARCH_TIME_S)
+
+
+def search(query):
+    global model, index
+
+    # TODO: pass in param string directly?
+    query_embedding = model.encode(query)
+    distances, faiss_ids = index.search("embeddings", query_embedding, 20)
+
+    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'
+
+    neighbors = [(distance, openalex_id) for distance, openalex_id in zip(distances, openalex_ids)]
+    request_str = f'https://api.openalex.org/works?filter={search_filter}&select={search_select}'
+
+    return neighbors, request_str
+
 with gr.Blocks() as demo:
     gr.Markdown('# abstracts-index')
     gr.Markdown(