Compute the similarity instead of the distance if applicable

app.py

@@ -118,6 +118,7 @@ def get_env_var[T, U](
 def get_model(
     model_name: str, params_dir: Path, trust_remote_code: bool
 ) -> tuple[bool, SentenceTransformer]:
+    # TODO: params["normalize"] for models like all-MiniLM-v6, which already normalize?
     with open(params_dir / "params.json", "r") as f:
         params: Params = json.load(f)
     return params["normalize"], SentenceTransformer(
@@ -173,7 +174,9 @@ def collapse_newlines(x: str) -> str:
     return x.replace("\r\n", " ").replace("\n", " ").replace("\r", " ")
 
 
-def format_response(neighbors: list[Work], distances: list[float]) -> str:
+def format_response(
+    neighbors: list[Work], distances: list[float], calculate_similarity: bool = False
+) -> str:
     result_string = ""
     for work, distance in zip(neighbors, distances):
         entry_string = "## "
@@ -218,7 +221,11 @@ def format_response(neighbors: list[Work], distances: list[float]) -> str:
             meta.append(("Cited-by count", str(work.citations)))
         if work.doi:
             meta.append(("DOI", work.doi.replace("https://doi.org/", "")))
-        meta.append(("Similarity", f"{distance:.2f}"))
+        if calculate_similarity:
+            # if query and result are unit vectors, the cosine sim is 1 - dist^2 / 2
+            meta.append(("Similarity", f"{1 - distance / 2:.2f}"))  # faiss gives dist^2
+        else:
+            meta.append(("Distance", f"{distance:.2f}"))
         entry_string += (" " * 4).join(": ".join(tup) for tup in meta)
 
         entry_string += "*\n"
@@ -264,7 +271,10 @@ def main():
         works_flat = execute_request(openalex_ids_flat, mailto)
         works = [list(batch) for batch in batched(works_flat, k)]
 
-        result_strings = [format_response(w, d) for w, d in zip(works, distances)]
+        result_strings = [
+            format_response(w, d, calculate_similarity=normalize)
+            for w, d in zip(works, distances)
+        ]
 
         return (result_strings, )