Create main function, add generic get_env_var function

app.py

@@ -6,16 +6,12 @@ from sentence_transformers import SentenceTransformer
 import faiss
 import gradio as gr
 from datasets import Dataset
-from typing import TypedDict, Self, Any
+from typing import TypedDict, Self, Any, Callable
 import json
 from pathlib import Path
 from dataclasses import dataclass
 from itertools import batched, chain
-
-MODEL_NAME = "all-MiniLM-L6-v2"  # TODO: make configurable
-DIR = Path("index")
-SEARCH_TIME_S = 1  # TODO: optimize
-K = 20
+import os
 
 
 class IndexParameters(TypedDict):
@@ -98,6 +94,17 @@ class Work:
         return " ".join(word for word in abstract_words if word is not None)
 
 
+def get_env_var[T, U](
+    key: str, type_: Callable[[str], T] = str, default: U = None
+) -> T | U:
+    var = os.getenv(key)
+    if var is not None:
+        var = type_(var)
+    else:
+        var = default
+    return var
+
+
 def get_model(model_name: str, device: str) -> SentenceTransformer:
     return SentenceTransformer(model_name, device=device)
 
@@ -107,6 +114,7 @@ def get_index(dir: Path, search_time_s: float) -> Dataset:
     index.load_faiss_index("embeddings", dir / "index.faiss", None)
     faiss_index: faiss.Index = index.get_index("embeddings").faiss_index  # type: ignore
 
+    # TODO: search for what minimized distance from utopia point
     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]
@@ -197,51 +205,62 @@ def format_response(neighbors: list[Work], distances: list[float]) -> str:
     return result_string
 
 
-model = get_model(MODEL_NAME, "cpu")
-index = get_index(DIR, SEARCH_TIME_S)
-
-
-# function signature: (expanded tuple of input batches) -> tuple of output batches
-def search(query: list[str]) -> tuple[list[str]]:
-    query_embedding = model.encode(query)
-    distances, faiss_ids = index.search_batch("embeddings", query_embedding, K)
-
-    faiss_ids_flat = list(chain(*faiss_ids))
-    openalex_ids_flat = index[faiss_ids_flat]["idxs"]
-    works_flat = execute_request(openalex_ids_flat)
-    works = [list(batch) for batch in batched(works_flat, K)]
-
-    result_strings = [format_response(w, d) for w, d in zip(works, distances)]
-
-    return (result_strings, )
+def main():
+    # TODO: figure out some better defaults?
+    model_name = get_env_var("MODEL_NAME", default="all-MiniLM-L6-v2")
+    dir = get_env_var("DIR", Path, default=Path("index"))
+    search_time_s = get_env_var("SEARCH_TIME_S", int, default=1)
+    k = get_env_var("K", int, default=20)
+
+    model = get_model(model_name, "cpu")
+    index = get_index(dir, search_time_s)
+
+    # function signature: (expanded tuple of input batches) -> tuple of output batches
+    def search(query: list[str]) -> tuple[list[str]]:
+        query_embedding = model.encode(query)
+        distances, faiss_ids = index.search_batch("embeddings", query_embedding, k)
+
+        faiss_ids_flat = list(chain(*faiss_ids))
+        openalex_ids_flat = index[faiss_ids_flat]["idxs"]
+        works_flat = execute_request(openalex_ids_flat)
+        works = [list(batch) for batch in batched(works_flat, k)]
+
+        result_strings = [format_response(w, d) for w, d in zip(works, distances)]
+
+        return (result_strings, )
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# abstracts-index")
+        gr.Markdown(
+            "Explore 95 million academic publications selected from the "
+            "[OpenAlex](https://openalex.org) dataset. This project is an index of the "
+            "embeddings generated from their titles and abstracts. The embeddings were "
+            f"generated using the {model_name} model provided by the "
+            "[sentence-transformers](https://www.sbert.net/) module, and the index was "
+            "built using the [faiss](https://github.com/facebookresearch/faiss) "
+            "module. The build scripts and more information available at the main repo "
+            "[abstracts-search](https://github.com/colonelwatch/abstracts-search) on "
+            "Github."
+        )
 
+        query = gr.Textbox(
+            lines=1, placeholder="Enter your query here", show_label=False
+        )
+        btn = gr.Button("Search")
+        results = gr.Markdown(
+            latex_delimiters=[
+                {"left": "$$", "right": "$$", "display": False},
+                {"left": "$", "right": "$", "display": False},
+            ],
+            container=True,
+        )
 
-with gr.Blocks() as demo:
-    gr.Markdown("# abstracts-index")
-    gr.Markdown(
-        "Explore 95 million academic publications selected from the "
-        "[OpenAlex](https://openalex.org) dataset. This project is an index of the "
-        "embeddings generated from their titles and abstracts. The embeddings were "
-        "generated using the `all-MiniLM-L6-v2` model provided by the "
-        "[sentence-transformers](https://www.sbert.net/) module, and the index was "
-        "built using the [faiss](https://github.com/facebookresearch/faiss) module. "
-        "The build scripts and more information available at the main repo "
-        "[abstracts-search](https://github.com/colonelwatch/abstracts-search) on "
-        "Github."
-    )
+        query.submit(search, inputs=[query], outputs=[results], batch=True)
+        btn.click(search, inputs=[query], outputs=[results], batch=True)
 
-    query = gr.Textbox(lines=1, placeholder="Enter your query here", show_label=False)
-    btn = gr.Button("Search")
-    results = gr.Markdown(
-        latex_delimiters=[
-            {"left": "$$", "right": "$$", "display": False},
-            {"left": "$", "right": "$", "display": False},
-        ],
-        container=True,
-    )
+    demo.queue()
+    demo.launch()
 
-    query.submit(search, inputs=[query], outputs=[results], batch=True)
-    btn.click(search, inputs=[query], outputs=[results], batch=True)
 
-demo.queue()
-demo.launch()
+if __name__ == "__main__":
+    main()