Normalize column names

app.py

@@ -128,8 +128,8 @@ def get_model(
 
 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
+    index.load_faiss_index("embedding", dir / "index.faiss", None)
+    faiss_index: faiss.Index = index.get_index("embedding").faiss_index  # type: ignore
 
     with open(dir / "params.json", "r") as f:
         params: Params = json.load(f)
@@ -248,10 +248,10 @@ def main():
         query_embedding = model.encode(
             query, prompt_name, normalize_embeddings=normalize
         )
-        distances, faiss_ids = index.search_batch("embeddings", query_embedding, k)
+        distances, faiss_ids = index.search_batch("embedding", query_embedding, k)
 
         faiss_ids_flat = list(chain(*faiss_ids))
-        openalex_ids_flat = index[faiss_ids_flat]["idxs"]
+        openalex_ids_flat = index[faiss_ids_flat]["id"]
         works_flat = execute_request(openalex_ids_flat, mailto)
         works = [list(batch) for batch in batched(works_flat, k)]