Spaces:
Running
Running
Delete server_embed.py
Browse files- server_embed.py +0 -123
server_embed.py
DELETED
|
@@ -1,123 +0,0 @@
|
|
| 1 |
-
import os
|
| 2 |
-
import functools
|
| 3 |
-
import numpy as np
|
| 4 |
-
|
| 5 |
-
import fastapi
|
| 6 |
-
import voyageai
|
| 7 |
-
vo = voyageai.Client(os.environ['VOYAGE_API_KEY'])
|
| 8 |
-
|
| 9 |
-
theorems = [
|
| 10 |
-
theorem.split('--ast--\n')
|
| 11 |
-
for theorem in open('ast.txt').read().split('----ast----\n')[1:]
|
| 12 |
-
]
|
| 13 |
-
embed = np.load('embed_arr.npy')
|
| 14 |
-
|
| 15 |
-
app = fastapi.FastAPI()
|
| 16 |
-
|
| 17 |
-
top_n = 50
|
| 18 |
-
|
| 19 |
-
@functools.lru_cache(maxsize=2*1024*1024*1024//(80 + top_n*8))
|
| 20 |
-
def lru_search(query):
|
| 21 |
-
query_embed = vo.embed([query], model="voyage-large-2-instruct", input_type="query").embeddings[0]
|
| 22 |
-
similarities = np.dot(embed, query_embed)
|
| 23 |
-
|
| 24 |
-
top = np.argpartition(similarities, -top_n)[-top_n:]
|
| 25 |
-
top = top[np.argsort(similarities[top])][::-1]
|
| 26 |
-
sim = similarities[top]
|
| 27 |
-
return top, sim
|
| 28 |
-
|
| 29 |
-
@app.get("/search")
|
| 30 |
-
def search(query: str):
|
| 31 |
-
# if query is too long, throw an error
|
| 32 |
-
if len(query) > 80:
|
| 33 |
-
raise fastapi.HTTPException(status_code=400, detail="Query is too long")
|
| 34 |
-
top, sim = lru_search(query)
|
| 35 |
-
|
| 36 |
-
results = []
|
| 37 |
-
for i, similarity in zip(top, sim):
|
| 38 |
-
module, name, typ, comment = theorems[i]
|
| 39 |
-
results.append({
|
| 40 |
-
"module": module,
|
| 41 |
-
"name": name,
|
| 42 |
-
"type": typ,
|
| 43 |
-
"comment": comment,
|
| 44 |
-
"similarity": similarity
|
| 45 |
-
})
|
| 46 |
-
return results
|
| 47 |
-
|
| 48 |
-
@app.get("/")
|
| 49 |
-
def read_root():
|
| 50 |
-
# HTML
|
| 51 |
-
return fastapi.responses.HTMLResponse(
|
| 52 |
-
content="""
|
| 53 |
-
<html>
|
| 54 |
-
<head>
|
| 55 |
-
<title>Search</title>
|
| 56 |
-
<!-- tailwind css -->
|
| 57 |
-
<script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
|
| 58 |
-
<link href="https://cdn.jsdelivr.net/npm/tailwindcss@2.2.19/dist/tailwind.min.css" rel="stylesheet">
|
| 59 |
-
</head>
|
| 60 |
-
<body class="mx-24 my-8 text-lg flex flex-col">
|
| 61 |
-
<h1 class="text-3xl">Search Mathlib</h1>
|
| 62 |
-
<form onsubmit="search(); return false;" class="flex gap-2 mt-2 mb-3">
|
| 63 |
-
<input autofocus type="text" name="query" class="grow border p-2 w-full rounded-lg focus:outline-none border-gray-300 focus:ring-1 focus:ring-gray-100">
|
| 64 |
-
<button type="submit" class="bg-green-700 text-white p-2 w-24 rounded-lg">Search</button>
|
| 65 |
-
</form>
|
| 66 |
-
<ul id="results" class="list-disc">
|
| 67 |
-
</ul>
|
| 68 |
-
<script>
|
| 69 |
-
function removeFirstWord(str) {
|
| 70 |
-
return str.split(' ').slice(1).join(' ')
|
| 71 |
-
}
|
| 72 |
-
// https://leanprover-community.github.io/mathlib4_docs/Mathlib/Tactic.Widget.SelectPanelUtils.html#def%20mkSelectionPanelRPC.match_1
|
| 73 |
-
|
| 74 |
-
async function search() {
|
| 75 |
-
const query = document.querySelector('input[name="query"]').value
|
| 76 |
-
const response = await fetch(`/search?query=${query}`)
|
| 77 |
-
if (!response.ok) {
|
| 78 |
-
alert('Error: ' + (await response.json())['detail'])
|
| 79 |
-
return
|
| 80 |
-
}
|
| 81 |
-
const results = await response.json()
|
| 82 |
-
|
| 83 |
-
const resultsDiv = document.querySelector('#results')
|
| 84 |
-
resultsDiv.innerHTML = ''
|
| 85 |
-
|
| 86 |
-
for (const result of results) {
|
| 87 |
-
const resultDiv = document.createElement('li')
|
| 88 |
-
const link = "https://leanprover-community.github.io/mathlib4_docs/" + removeFirstWord(result.module).replaceAll('.', '/') + ".html#" + removeFirstWord(result.name)
|
| 89 |
-
resultDiv.innerHTML = `
|
| 90 |
-
<div class="flex gap-1">
|
| 91 |
-
<a href="${link}"
|
| 92 |
-
class="text-xl mr-4 text-green-800 hover:text-green-600
|
| 93 |
-
">${result.name}</a>
|
| 94 |
-
<div class="flex-grow"></div>
|
| 95 |
-
<p class="text-gray-800">Similarity: ${result.similarity.toFixed(3)}</p>
|
| 96 |
-
</div>
|
| 97 |
-
<div class="flex flex-col">
|
| 98 |
-
<p class="font-mono">: ${result.type}</p>
|
| 99 |
-
<p class="text-gray-800 mt-1">${result.module}</p>
|
| 100 |
-
<p id="comment"></p>
|
| 101 |
-
</div>
|
| 102 |
-
|
| 103 |
-
<div class="border my-1 border-gray-200"></div>
|
| 104 |
-
`
|
| 105 |
-
if (result.comment.trim() !== '') {
|
| 106 |
-
const comment = resultDiv.querySelector('#comment')
|
| 107 |
-
comment.innerHTML = marked.marked(`/-- ${result.comment.trim()} -/`)
|
| 108 |
-
comment.classList.add('bg-gray-50', 'text-gray-900', 'px-2', 'py-1', 'mt-0.5', 'rounded-lg')
|
| 109 |
-
}
|
| 110 |
-
resultsDiv.appendChild(resultDiv)
|
| 111 |
-
}
|
| 112 |
-
}
|
| 113 |
-
</script>
|
| 114 |
-
</body>
|
| 115 |
-
</html>
|
| 116 |
-
"""
|
| 117 |
-
)
|
| 118 |
-
|
| 119 |
-
'''
|
| 120 |
-
to run this with auto-reload, use the following command:
|
| 121 |
-
cd data
|
| 122 |
-
uvicorn test_embed:app --reload
|
| 123 |
-
'''
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|