Spaces:
Running
Running
| 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 | |
| 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 | |
| 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 | |
| def read_root(): | |
| # HTML | |
| return fastapi.responses.HTMLResponse( | |
| content=""" | |
| <html> | |
| <head> | |
| <title>Search Mathlib</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="overflow-scroll mx-auto w-10/12 my-8 text-lg flex flex-col"> | |
| <h1 class="text-3xl mx-0.5">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> | |
| <!-- The similarity info is only shown on larger screens --> | |
| <p class="text-gray-800 hidden md:block">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 server:app --reload | |
| ''' | |