import os import functools import numpy as np import fastapi import voyageai vo = voyageai.Client(os.environ['VOYAGE_API_KEY']) theorems = [ theorem.split('--thm--\n') for theorem in open('theorems.txt').read().split('----thm----\n')[1:] ] embed = np.load('embed_arr.npy') app = fastapi.FastAPI() top_n = 50 @functools.lru_cache(maxsize=2*1024*1024*1024//(400 + top_n*8)) def lru_search(query): query_embed = vo.embed([query], model="voyage-large-2-instruct", input_type="query").embeddings[0] similarities = np.dot(embed, query_embed) top = np.argpartition(similarities, -top_n)[-top_n:] top = top[np.argsort(similarities[top])][::-1] sim = similarities[top] return top, sim @app.get("/search") def search(query: str): # if query is too long, throw an error if len(query) > 400: raise fastapi.HTTPException(status_code=400, detail="Query is too long") top, sim = lru_search(query) results = [] for i, similarity in zip(top, sim): module, name, typ, comment = theorems[i] results.append({ "module": module, "name": name, "type": typ, "comment": comment, "similarity": similarity }) return results @app.get("/") def read_root(): # HTML return fastapi.responses.HTMLResponse( content=""" Search Mathlib

Search Mathlib

You can search for any object declared in Mathlib and Lean, including theorems, functions, types, syntax, tactics, and metaprogramming functions.

The neural network sees the name, type, file location, and comment of a declaration. Your search terms should match these as closely as possible.

You can find README.md and the code I used in the Files tab. Thank you!

""" ) ''' to run this with auto-reload, use the following command: cd data uvicorn server:app --reload '''