Change the vector database used and embed the embeddings within the program

.gitattributes

@@ -33,3 +33,5 @@ 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
+.kanjidb/**/* filter=lfs diff=lfs merge=lfs -text
+*.lance filter=lfs diff=lfs merge=lfs -text

.kanjidb/kanji.lance/_latest.manifest

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d870cc02d1237f025526565eba9054fdecb7e5f29f96230707b6418a1594f8e7
+size 279

.kanjidb/kanji.lance/_transactions/0-199d3b3b-6378-4c03-9465-cad215d61adf.txn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb34cb00ef08393e705a75aa78c4068c8da6c5385fc118afd506f18e5c956594
+size 165

.kanjidb/kanji.lance/_transactions/1-a2c30a19-a8de-47a2-ae7b-08ce8f374e71.txn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9963f78e3fbb2f07e662ecb999ddc0e6f276649ae8cc65d4a9159576fbf79e17
+size 100

.kanjidb/kanji.lance/_versions/1.manifest

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1abfe73ec38efbad945e6146850736e4572d7bd364b4361c2c4c595be35cf9cc
+size 222

.kanjidb/kanji.lance/_versions/2.manifest

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d870cc02d1237f025526565eba9054fdecb7e5f29f96230707b6418a1594f8e7
+size 279

.kanjidb/kanji.lance/data/866ef445-f6ad-45be-8ff1-620646e5f41f.lance

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:41f3f57790cfa577b99f2283d5a470914aed2082449f42fcb37d9b62634b4c6f
+size 238451553

README.md

@@ -4,9 +4,9 @@ emoji: 👀
 colorFrom: purple
 colorTo: gray
 sdk: gradio
-sdk_version: 4.18.0
+sdk_version: 4.31.4
 app_file: app.py
-pinned: false
+pinned: true
 license: mit
 ---
 

app.py

@@ -6,16 +6,15 @@ from config import (
     article,
 )
 from encode import get_embeddings
-from database import search_vector, format_search_results
+from database import search_vector
 
 def search_images(values):
     image = Image.new("RGBA", values["composite"].size, (255, 255, 255, 255))
     image.paste(values["composite"], mask=values["composite"])
     embedding = get_embeddings([image])[0]
     results = search_vector(embedding, limit=100)
-    formatted = format_search_results(results)
 
-    _deduplicated = '\t'.join(dict.fromkeys(result.kanji for result in formatted))
+    _deduplicated = '\t'.join(dict.fromkeys(result.kanji for result in results))
     # TODO Format the results better
     # Huge boxes using the right font for each of them?
 
@@ -27,7 +26,7 @@ input_image = gr.ImageEditor(
     show_label=False,
     type="pil",
     brush=gr.Brush(
-        default_size=3,
+        default_size=16,
         color_mode="fixed",
         colors=["#000000", "#ffffff"],
     ),

config.py

@@ -1,7 +1,6 @@
 import os
 
-qdrant_location = os.getenv('QDRANT_URL', "localhost")
-qdrant_api_key = os.getenv('QDRANT_API_KEY')
+lancedb_location = os.getenv('DATABASE_FILE', ".kanjidb")
 
 description = """This is a Kanji image search demo. Draw or upload an image of an individual Kanji character."""
 
@@ -16,7 +15,7 @@ The results is sorted by estimated distance from the input, but will rarely give
 
 ### About this project
 
-It uses the "kha-white/manga-ocr-base" Vision Transformer Encoder model to create embeddings, then uses a vector database (qdrant) to find similar characters.
+It uses the "kha-white/manga-ocr-base" Vision Transformer Encoder model to create embeddings, then uses a vector database (lancedb) to find similar characters.
 
 You can find the code used to create the embeddings as well as more information in https://github.com/etrotta/kanji_lookup
 

database.py

@@ -1,39 +1,27 @@
-import dataclasses
 import torch
-from qdrant_client import QdrantClient, models
+import lancedb
+from lancedb.pydantic import LanceModel
+import pydantic
+# import time
 
-from config import qdrant_location, qdrant_api_key
+from config import lancedb_location
 
-qdrant = QdrantClient(
-    qdrant_location,
-    api_key=qdrant_api_key,
-    port=443,
-    timeout=30,
-)
+db = lancedb.connect(lancedb_location)
+table = db.open_table("kanji")
 
-def search_vector(query_vector: torch.Tensor, limit: int=20) -> list[models.ScoredPoint]:
-    hits = qdrant.search(
-        collection_name="kanji",
-        # query_vector=query_vector,
-        query_vector=query_vector.numpy(),
-        limit=limit,
-        with_payload=True,
-    )
-    return hits
-
-@dataclasses.dataclass
-class SearchResult:
+class SearchResult(LanceModel):
     kanji: str
-    font: str
-    score: float
+    distance: float = pydantic.Field(validation_alias=pydantic.AliasChoices('distance', '_distance'))
 
-def format_search_results(hits: list[models.ScoredPoint]) -> list[SearchResult]:
-    formatted = []
-    for point in hits:
-        kanji, font = point.payload["kanji"], point.payload["font"]
-        formatted.append(SearchResult(
-            kanji = kanji,
-            font = font,
-            score = point.score,
-        ))
-    return formatted
+def search_vector(query_vector: torch.Tensor, limit: int=20) -> list[SearchResult]:
+    # start = time.perf_counter()
+    results = (
+        table
+        .search(query_vector.numpy(), vector_column_name="vector", query_type="vector")
+        .limit(limit)
+        # .to_pydantic(SearchResult)   # type: ignore
+        .to_list()
+    )
+    # end = time.perf_counter()
+    # print(f"Searched in {end - start:.3f}")
+    return [SearchResult.model_validate(result) for result in results]