dx2102 commited on
Commit
b65779c
1 Parent(s): 87b632e

Upload 2 files

Browse files
Files changed (3) hide show
  1. .gitattributes +1 -0
  2. server.py +123 -0
  3. theorems.txt +3 -0
.gitattributes CHANGED
@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
  ast.txt filter=lfs diff=lfs merge=lfs -text
 
 
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
  ast.txt filter=lfs diff=lfs merge=lfs -text
37
+ theorems.txt filter=lfs diff=lfs merge=lfs -text
server.py ADDED
@@ -0,0 +1,123 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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('--thm--\n')
11
+ for theorem in open('theorems.txt').read().split('----thm----\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 server:app --reload
123
+ '''
theorems.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ec39405cd29f0cdb25dde5079567825ca9e40f07cde8772056282272cb38ef33
3
+ size 99276292