.

indexer.py

@@ -133,7 +133,10 @@ def main() -> None:
     model.eval()
 
     client = chromadb.PersistentClient(path=str(args.chroma_path))
-    collection = client.get_or_create_collection(name=args.collection)
+    collection = client.get_or_create_collection(
+        name=args.collection,
+        metadata={"hnsw:space": "cosine"},
+    )
 
     total_images = len(image_paths)
     progress = tqdm(total=total_images, desc="Indexing images", unit="img")

inference.py

@@ -21,7 +21,7 @@ def parse_args() -> argparse.Namespace:
         default=[
             "레인보우 스타",
             "블랙과 흰색의 별 모양 무기",
-            "핑크색 날개",
+            "흰색 티셔츠",
             "파란색 모자",
             "관련 없는 이미지",
         ],
@@ -30,6 +30,15 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
+def resolve_adapter_path(adapter_path: Path) -> Path:
+    if (adapter_path / "adapter_config.json").exists():
+        return adapter_path
+    candidate = adapter_path / "best_model"
+    if (candidate / "adapter_config.json").exists():
+        return candidate
+    return adapter_path
+
+
 def main() -> None:
     args = parse_args()
 
@@ -38,7 +47,8 @@ def main() -> None:
 
     print("Loading model...")
     base_model = SiglipModel.from_pretrained(args.model_id)
-    model = PeftModel.from_pretrained(base_model, args.adapter_path)
+    adapter_path = resolve_adapter_path(Path(args.adapter_path))
+    model = PeftModel.from_pretrained(base_model, str(adapter_path))
     processor = SiglipProcessor.from_pretrained(args.model_id)
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -59,13 +69,12 @@ def main() -> None:
         image_embeds = image_embeds / image_embeds.norm(dim=-1, keepdim=True)
         text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True)
 
-        logits = image_embeds @ text_embeds.t()
-        logit_scale = model.logit_scale.exp()
-        logits = logits * logit_scale
-        probs = logits.softmax(dim=1)
+        similarities = image_embeds @ text_embeds.t()
 
-    for text, prob in zip(args.candidates, probs[0]):
-        print(f"{text}: {prob.item() * 100:.2f}%")
+    for text, similarity in zip(args.candidates, similarities[0]):
+        sim_value = similarity.item()
+        distance = 1.0 - sim_value
+        print(f"{text} Similarity: {sim_value:.4f} | Distance: {distance:.4f}")
 
 
 if __name__ == "__main__":

main.py

@@ -49,7 +49,10 @@ async def lifespan(app: FastAPI):
     model.eval()
 
     client = chromadb.PersistentClient(path="chroma_db")
-    collection = client.get_or_create_collection(name="maple_items")
+    collection = client.get_or_create_collection(
+        name="maple_items",
+        metadata={"hnsw:space": "cosine"},
+    )
 
     app.state.device = device
     app.state.model = model
@@ -97,7 +100,7 @@ def search(payload: SearchRequest) -> Dict[str, Any]:
     results = collection.query(
         query_embeddings=[query_embedding],
         n_results=payload.k,
-        include=["distances", "metadatas", "ids"],
+        include=["distances", "metadatas"],
     )
 
     ids: List[str] = results.get("ids", [[]])[0]
@@ -110,11 +113,13 @@ def search(payload: SearchRequest) -> Dict[str, Any]:
         if metadata:
             filepath = metadata.get("filepath", "")
         image_url = f"/static/images/{filepath}" if filepath else ""
+        similarity = max(0.0, 1.0 - distance) if distance is not None else 0.0
         response_items.append(
             {
                 "id": item_id,
                 "filepath": filepath,
                 "distance": distance,
+                "similarity": similarity,
                 "image_url": image_url,
             }
         )

requirements.txt

@@ -11,3 +11,6 @@ sentencepiece>=0.1.99
 protobuf>=4.21
 peft>=0.11
 scikit-learn>=1.3
+chromadb>=0.5
+fastapi>=0.110
+uvicorn>=0.29