Spaces:

dx2102
/

search-mathlib

Running

App Files Files Community

dx2102 commited on Jul 29, 2024

Commit

e8eca0b

verified ·

1 Parent(s): a98a2b2

Upload 3 files

Browse files

Files changed (4) hide show

.gitattributes +1 -0
ast.txt +3 -0
embed_arr.npy +3 -0
server_embed.py +123 -0

.gitattributes CHANGED Viewed

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text

 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+ast.txt filter=lfs diff=lfs merge=lfs -text

ast.txt ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:32075dd1929d18fe6cad3e7d9f0ad9346dea43530bcc374210316b03213e6c23
+size 99276292

embed_arr.npy ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:6b370e9bfaa79c0bd4f7204ef4a2eb6889404726b0c71244a9aabf1888c648d4
+size 2293915776

server_embed.py ADDED Viewed

	@@ -0,0 +1,123 @@

+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="""
+<html>
+<head>
+    <title>Search</title>
+    <!-- tailwind css -->
+    <script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
+    <link href="https://cdn.jsdelivr.net/npm/tailwindcss@2.2.19/dist/tailwind.min.css" rel="stylesheet">
+</head>
+<body class="mx-24 my-8 text-lg flex flex-col">
+    <h1 class="text-3xl">Search Mathlib</h1>
+    <form onsubmit="search(); return false;" class="flex gap-2 mt-2 mb-3">
+        <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">
+        <button type="submit" class="bg-green-700 text-white p-2 w-24 rounded-lg">Search</button>
+    </form>
+    <ul id="results" class="list-disc">
+    </ul>
+    <script>
+        function removeFirstWord(str) {
+            return str.split(' ').slice(1).join(' ')
+        }
+        // https://leanprover-community.github.io/mathlib4_docs/Mathlib/Tactic.Widget.SelectPanelUtils.html#def%20mkSelectionPanelRPC.match_1
+        async function search() {
+            const query = document.querySelector('input[name="query"]').value
+            const response = await fetch(`/search?query=${query}`)
+            if (!response.ok) {
+                alert('Error: ' + (await response.json())['detail'])
+                return
+            }
+            const results = await response.json()
+            const resultsDiv = document.querySelector('#results')
+            resultsDiv.innerHTML = ''
+            for (const result of results) {
+                const resultDiv = document.createElement('li')
+                const link = "https://leanprover-community.github.io/mathlib4_docs/" + removeFirstWord(result.module).replaceAll('.', '/') + ".html#" + removeFirstWord(result.name)
+                resultDiv.innerHTML = `
+                    <div class="flex gap-1">
+                        <a href="${link}"
+                        class="text-xl mr-4 text-green-800 hover:text-green-600
+                          ">${result.name}</a>
+                        <div class="flex-grow"></div>
+                        <p class="text-gray-800">Similarity: ${result.similarity.toFixed(3)}</p>
+                    </div>
+                    <div class="flex flex-col">
+                        <p class="font-mono">: ${result.type}</p>
+                        <p class="text-gray-800 mt-1">${result.module}</p>
+                        <p id="comment"></p>
+                    </div>
+                    <div class="border my-1 border-gray-200"></div>
+                `
+                if (result.comment.trim() !== '') {
+                    const comment = resultDiv.querySelector('#comment')
+                    comment.innerHTML = marked.marked(`/-- ${result.comment.trim()} -/`)
+                    comment.classList.add('bg-gray-50', 'text-gray-900', 'px-2', 'py-1', 'mt-0.5', 'rounded-lg')
+                }
+                resultsDiv.appendChild(resultDiv)
+            }
+        }
+    </script>
+</body>
+</html>
+    """
+    )
+'''
+to run this with auto-reload, use the following command:
+cd data
+uvicorn test_embed:app --reload
+'''