app

.gitignore

@@ -0,0 +1,2 @@
+.env
+data

app.py

@@ -0,0 +1,141 @@
+import os
+from typing import ClassVar
+
+# import dotenv
+import gradio as gr
+import lancedb
+import srsly
+from huggingface_hub import snapshot_download
+from lancedb.embeddings.base import TextEmbeddingFunction
+from lancedb.embeddings.registry import register
+from lancedb.pydantic import LanceModel, Vector
+from lancedb.rerankers import CohereReranker, ColbertReranker
+from lancedb.util import attempt_import_or_raise
+
+# dotenv.load_dotenv()
+
+
+@register("coherev3")
+class CohereEmbeddingFunction_2(TextEmbeddingFunction):
+    name: str = "embed-english-v3.0"
+    client: ClassVar = None
+
+    def ndims(self):
+        return 768
+
+    def generate_embeddings(self, texts):
+        """
+        Get the embeddings for the given texts
+
+        Parameters
+        ----------
+        texts: list[str] or np.ndarray (of str)
+            The texts to embed
+        """
+        # TODO retry, rate limit, token limit
+        self._init_client()
+        rs = CohereEmbeddingFunction_2.client.embed(
+            texts=texts, model=self.name, input_type="search_document"
+        )
+
+        return [emb for emb in rs.embeddings]
+
+    def _init_client(self):
+        cohere = attempt_import_or_raise("cohere")
+        if CohereEmbeddingFunction_2.client is None:
+            CohereEmbeddingFunction_2.client = cohere.Client(
+                os.environ["COHERE_API_KEY"]
+            )
+
+
+COHERE_EMBEDDER = CohereEmbeddingFunction_2.create()
+
+
+class ArxivModel(LanceModel):
+    text: str = COHERE_EMBEDDER.SourceField()
+    vector: Vector(1024) = COHERE_EMBEDDER.VectorField()
+    title: str
+    paper_title: str
+    content_type: str
+    arxiv_id: str
+
+
+def download_data():
+    snapshot_download(
+        repo_id="rbiswasfc/zotero_db",
+        repo_type="dataset",
+        local_dir="./data",
+        token=os.environ["HF_TOKEN"],
+    )
+    print("Data downloaded!")
+
+
+download_data()
+
+VERSION = "0.0.0a"
+DB = lancedb.connect("./data/.lancedb_zotero_v0")
+ID_TO_ABSTRACT = srsly.read_json("./data/id_to_abstract.json")
+RERANKERS = {"colbert": ColbertReranker(), "cohere": CohereReranker()}
+TBL = DB.open_table("arxiv_zotero_v0")
+
+
+def _format_results(arxiv_refs):
+    results = []
+    for arx_id, paper_title in arxiv_refs.items():
+        abstract = ID_TO_ABSTRACT.get(arx_id, "")
+        # these are all ugly hacks because the data preprocessing is poor. to be fixed v soon.
+        if "Abstract\n\n" in abstract:
+            abstract = abstract.split("Abstract\n\n")[-1]
+        if paper_title in abstract:
+            abstract = abstract.split(paper_title)[-1]
+        if abstract.startswith("\n"):
+            abstract = abstract[1:]
+        if "\n\n" in abstract[:20]:
+            abstract = "\n\n".join(abstract.split("\n\n")[1:])
+        result = {
+            "title": paper_title,
+            "url": f"https://arxiv.org/abs/{arx_id}",
+            "abstract": abstract,
+        }
+        results.append(result)
+
+    return results
+
+
+def query_db(query: str, k: int = 10, reranker: str = "cohere"):
+    raw_results = TBL.search(query, query_type="hybrid").limit(k)
+    if reranker is not None:
+        ranked_results = raw_results.rerank(reranker=RERANKERS[reranker])
+    else:
+        ranked_results = raw_results
+
+    ranked_results = ranked_results.to_pandas()
+    top_results = ranked_results.groupby("arxiv_id").agg({"_relevance_score": "sum"})
+    top_results = top_results.sort_values(by="_relevance_score", ascending=False).head(
+        3
+    )
+    top_results_dict = {
+        row["arxiv_id"]: row["paper_title"]
+        for index, row in ranked_results.iterrows()
+        if row["arxiv_id"] in top_results.index
+    }
+
+    final_results = _format_results(top_results_dict)
+    return final_results
+
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        query = gr.Textbox(label="Query", placeholder="Enter your query...")
+        submit_btn = gr.Button("Submit")
+    output = gr.JSON(label="Search Results")
+
+    # # callback ---
+    submit_btn.click(
+        fn=query_db,
+        inputs=query,
+        outputs=output,
+    )
+
+
+demo.launch()

requirements.txt

@@ -0,0 +1,9 @@
+transformers
+torch
+lancedb
+srsly
+cohere
+python-dotenv
+tantivy
+beautifulsoup4
+retry