import os
import functools
import numpy as np
import fastapi
import voyageai
vo = voyageai.Client(os.environ['VOYAGE_API_KEY'])
theorems = [
theorem.split('--ast--\n')
for theorem in open('ast.txt').read().split('----ast----\n')[1:]
]
embed = np.load('embed_arr.npy')
app = fastapi.FastAPI()
top_n = 50
@functools.lru_cache(maxsize=2*1024*1024*1024//(80 + 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) > 80:
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
Search Mathlib
"""
)
'''
to run this with auto-reload, use the following command:
cd data
uvicorn test_embed:app --reload
'''