Implement text search

Author: qbonnard

Makefile

@@ -53,14 +53,16 @@ test:
 # The echo commands informally compare the actual to the expected result.
 API_URL=localhost:4213
 integration-test:
-	curl -X POST ${API_URL}/models/vit_b32/urls | cut -c 1-100
+	curl -s -X POST ${API_URL}/models/vit_b32/urls | cut -c 1-100
 	@echo =?=[]
 	curl -H "Content-Type: application/json" -d '{"url": "https://iiif.itatti.harvard.edu/iiif/2/yashiro!letters-jp!letter_001.pdf/full/full/0/default.jpg"}' ${API_URL}/models/vit_b32/add
 	@echo =?=
-	curl -X POST ${API_URL}/models/vit_b32/urls | cut -c 1-100
+	curl -s -X POST ${API_URL}/models/vit_b32/urls | cut -c 1-100
 	@echo =?=[url]
 	curl -H "Content-Type: application/json" -d '{"url": "https://iiif.itatti.harvard.edu/iiif/2/yashiro!letters-jp!letter_001.pdf/full/full/0/default.jpg"}' ${API_URL}/models/vit_b32/search
 	@echo =?=[result]
+	curl -H "Content-Type: application/json" -d '{"text": "a black and white text in japanese"}' ${API_URL}/models/vit_b32/search
+	@echo =?=[result]
 	curl -H "Content-Type: application/json" -d '{"url": "https://iiif.itatti.harvard.edu/iiif/2/yashiro!letters-jp!letter_001.pdf/full/full/0/default.jpg", "other": "https://iiif.itatti.harvard.edu/iiif/2/yashiro!letters-jp!letter_001.pdf/full/full/0/default.jpg"}' ${API_URL}/models/vit_b32/compare
 	@echo =?=100
 	curl ${API_URL}/models/vit_b32/count

api/commands.py

@@ -25,7 +25,10 @@ def insert_images(
 
     new_urls = [url for url in urls if url not in existing_urls]
     image_embeddings = (
-        [embeddings[model_name].extract(load_image_from_url(url)) for url in new_urls]
+        [
+            embeddings[model_name].get_image_embedding(load_image_from_url(url))
+            for url in new_urls
+        ]
         if image_embeddings is None
         else [
             embedding
@@ -53,8 +56,7 @@ def similarity_score(distance):
     return 100 * (1 - distance / 2)
 
 
-def search(model_name, url, limit=10):
-    embedding = embeddings[model_name].extract(load_image_from_url(url))
+def search_by_embedding(model_name, embedding, limit=10):
     search_results = collections[model_name].search(
         data=[embedding],
         anns_field="embedding",
@@ -79,11 +81,21 @@ def search(model_name, url, limit=10):
     ]
 
 
+def search_by_url(model_name, url, limit=10):
+    embedding = embeddings[model_name].get_image_embedding(load_image_from_url(url))
+    return search_by_embedding(model_name, embedding, limit)
+
+
+def search_by_text(model_name, text, limit=10):
+    embedding = embeddings[model_name].get_text_embedding(text)
+    return search_by_embedding(model_name, embedding, limit)
+
+
 def compare(model_name, url_left, url_right):
     # alternatively, we could first try to fetch the embeddings from milvus in
     # case their computation is significantly more expensive than a query
-    left = embeddings[model_name].extract(load_image_from_url(url_left))
-    right = embeddings[model_name].extract(load_image_from_url(url_right))
+    left = embeddings[model_name].get_image_embedding(load_image_from_url(url_left))
+    right = embeddings[model_name].get_image_embedding(load_image_from_url(url_right))
 
     # calc_distance() has been removed from milvus
     # it's a bit overkill anyway if we don't compare with vectors from the db
@@ -140,7 +152,8 @@ def remove_images(model_name, urls):
 
 commands = dict(
     insert_images=insert_images,
-    search=search,
+    search_by_url=search_by_url,
+    search_by_text=search_by_text,
     compare=compare,
     list_images=list_images,
     count=count,

api/embeddings.py

@@ -37,7 +37,7 @@ def get_text_embedding(self, text):
         text_embedding = self.model.get_text_features(**inputs)
         text_embedding /= text_embedding.norm(dim=-1, keepdim=True)
         text_embedding = text_embedding.tolist()
-        return text_embedding
+        return text_embedding[0]
 
     @torch.no_grad()
     def get_image_embedding(self, images):
@@ -46,10 +46,7 @@ def get_image_embedding(self, images):
         image_embedding = self.model.get_image_features(**inputs)
         image_embedding /= image_embedding.norm(dim=-1, keepdim=True)
         image_embedding = image_embedding.tolist()
-        return image_embedding
-
-    def extract(self, image):
-        return self.get_image_embedding(image)[0]
+        return image_embedding[0]
 
 
 embeddings = {

api/main.py

@@ -240,9 +240,9 @@ async def restore(model_name: ModelName, images: list[DatabaseEntry]):
 )
 async def search(model_name: ModelName, params: SearchParameters):
     if params.url:
-        return try_rpc("search", [model_name.value, params.url, params.limit])
+        return try_rpc("search_by_url", [model_name.value, params.url, params.limit])
     elif params.text:
-        return try_rpc("text_search", [model_name.value, params.text, params.limit])
+        return try_rpc("search_by_text", [model_name.value, params.text, params.limit])
     else:
         raise HTTPException(
             status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,

api/tests/test_commands.py

@@ -44,7 +44,21 @@ def test_crud(mock_model):
             "metadata": metadata,
             "url": TEST_URLS[0],
         }
-    ] == commands["search"](mock_model, TEST_URLS[1])
+    ] == commands["search_by_url"](mock_model, TEST_URLS[1])
+    assert [
+        {
+            "similarity": pytest.approx(29.19090986251831),
+            "metadata": metadata,
+            "url": TEST_URLS[0],
+        }
+    ] == commands["search_by_text"](mock_model, "a black and white text in japanese")
+    assert [
+        {
+            "similarity": pytest.approx(17.36249327659607),
+            "metadata": metadata,
+            "url": TEST_URLS[0],
+        }
+    ] == commands["search_by_text"](mock_model, "a cute colorful cat")
 
     commands["remove_images"](mock_model, [TEST_URLS[0]])
 
@@ -65,9 +79,9 @@ def test_search_more_results(mock_model):
         [None] * 1000,
         [[0] * 512] * 1000,
     )
-    results = commands["search"](mock_model, TEST_URLS[1])
+    results = commands["search_by_url"](mock_model, TEST_URLS[1])
     assert 10 == len(results)
-    results = commands["search"](mock_model, TEST_URLS[1], 100)
+    results = commands["search_by_url"](mock_model, TEST_URLS[1], 100)
     assert 100 == len(results)
     commands["remove_images"](mock_model, urls)
 

api/tests/test_embeddings.py

@@ -31,18 +31,26 @@ def squared_l2(v):
     return sum([x * x for x in v])
 
 
-def test_extract_vit_b32():
+def test_vit_b32_get_image_embedding():
     random.seed(2023)
     image_data = bytes([random.randint(0, 255) for _ in range(500 * 500 * 3)])
     image = Image.frombytes("RGB", (500, 500), image_data)
 
     b32 = embeddings["vit_b32"]
-    embedding = b32.extract(image)
+    embedding = b32.get_image_embedding(image)
     assert pytest.approx(1) == squared_l2(embedding)
     assert 512 == len(embedding)
     assert pytest.approx(-1.00633253128035) == sum(embedding)
 
 
+def test_vit_b32_get_text_embedding():
+    b32 = embeddings["vit_b32"]
+    embedding = b32.get_text_embedding("some random text")
+    assert pytest.approx(1) == squared_l2(embedding)
+    assert 512 == len(embedding)
+    assert pytest.approx(1.223632167381993) == sum(embedding)
+
+
 @functools.cache
 def batch_test_images():
     batch_test_ids = (
@@ -68,14 +76,14 @@ def batch_test_images():
 def test_individual_extract(benchmark):
     def extract_each():
         for image in batch_test_images():
-            embeddings["vit_b32"].extract(image)
+            embeddings["vit_b32"].get_image_embedding(image)
 
     benchmark(extract_each)
 
 
 @pytest.mark.benchmark
 def test_batch_extract(benchmark):
-    benchmark(embeddings["vit_b32"].extract, batch_test_images())
+    benchmark(embeddings["vit_b32"].get_image_embedding, batch_test_images())
 
 
 def test_unresized_image():
@@ -84,4 +92,4 @@ def test_unresized_image():
     url2 = "https://artresearch-iiif.s3.eu-west-1.amazonaws.com/marburg/gm1159076.jpg"
 
     images = [load_image_from_url(url) for url in [url0, url1, url2]]
-    [embeddings["vit_b32"].extract(image) for image in images]
+    [embeddings["vit_b32"].get_image_embedding(image) for image in images]

api/tests/test_main.py

@@ -180,6 +180,7 @@ def test_search_add_image_conflict(mock_rpc):
         },
     )
     assert response.status_code == 409
+    assert response.json() == {"detail": "Image already inserted"}
     mock_rpc.assert_called_once_with(
         "insert_images",
         [
@@ -192,7 +193,7 @@ def test_search_add_image_conflict(mock_rpc):
     )
 
 
-def test_search_success(mock_rpc):
+def test_search_by_url_success(mock_rpc):
     mock_rpc.return_value = [
         {
             "url": "http://example.com/image1.jpg",
@@ -223,11 +224,11 @@ def test_search_success(mock_rpc):
         },
     ]
     mock_rpc.assert_called_once_with(
-        "search", ["vit_b32", "http://example.com/query.jpg", 10]
+        "search_by_url", ["vit_b32", "http://example.com/query.jpg", 10]
     )
 
 
-def test_search_success_with_limit(mock_rpc):
+def test_search_by_url_success_with_limit(mock_rpc):
     mock_rpc.return_value = [
         {
             "url": "http://example.com/image1.jpg",
@@ -242,7 +243,7 @@ def test_search_success_with_limit(mock_rpc):
     assert response.status_code == 200
     assert len(response.json()) == 100
     mock_rpc.assert_called_once_with(
-        "search", ["vit_b32", "http://example.com/query.jpg", 100]
+        "search_by_url", ["vit_b32", "http://example.com/query.jpg", 100]
     )
 
 
@@ -255,20 +256,30 @@ def test_search_empty(mock_rpc):
     assert response.status_code == 200
     assert response.json() == []
     mock_rpc.assert_called_once_with(
-        "search", ["vit_b32", "http://example.com/query.jpg", 10]
+        "search_by_url", ["vit_b32", "http://example.com/query.jpg", 10]
     )
 
 
-def test_search_returns_422_when_invalid_url(mock_rpc):
+def test_search_by_url_returns_422_when_invalid_url(mock_rpc):
     response = client.post(
         "/models/vit_b32/search",
         json={"url": "not_a_url"},
     )
     assert response.status_code == 422
+    assert response.json() == {
+        "detail": [
+            {
+                "loc": ["body", "url"],
+                "msg": "invalid or missing URL scheme",
+                "type": "value_error.url.scheme",
+            }
+        ]
+    }
+
     mock_rpc.assert_not_called()
 
 
-def test_search_text_success(mock_rpc):
+def test_search_by_text_success(mock_rpc):
     mock_rpc.return_value = [
         {
             "url": "http://example.com/image1.jpg",
@@ -298,7 +309,24 @@ def test_search_text_success(mock_rpc):
             "similarity": 20,
         },
     ]
-    mock_rpc.assert_called_once_with("text_search", ["vit_b32", "cute cat", 10])
+    mock_rpc.assert_called_once_with("search_by_text", ["vit_b32", "cute cat", 10])
+
+
+def test_search_by_text_success_with_limit(mock_rpc):
+    mock_rpc.return_value = [
+        {
+            "url": "http://example.com/image1.jpg",
+            "metadata": {"tags": ["cat", "cute"]},
+            "similarity": 10,
+        }
+    ] * 100
+    response = client.post(
+        "/models/vit_b32/search",
+        json={"text": "some text", "limit": 100},
+    )
+    assert response.status_code == 200
+    assert len(response.json()) == 100
+    mock_rpc.assert_called_once_with("search_by_text", ["vit_b32", "some text", 100])
 
 
 def test_search_returns_422_whithout_query(mock_rpc):
@@ -319,7 +347,7 @@ def test_search_returns_500_when_rpc_error(mock_rpc):
     )
     assert response.status_code == 500
     mock_rpc.assert_called_once_with(
-        "search", ["vit_b32", "http://example.com/query.jpg", 10]
+        "search_by_url", ["vit_b32", "http://example.com/query.jpg", 10]
     )