bug fixing

pepkit · Dec 11, 2023 · 80e9c25 · 80e9c25
1 parent 304cbcc
commit 80e9c25
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/.gitignore b/.gitignore
@@ -140,4 +140,5 @@ postgres/
 .idea
 environment/local.env
 qdrant_storage/
-postgres-data/
+postgres-data/
+local_cache/
diff --git a/pephub/routers/api/v1/search.py b/pephub/routers/api/v1/search.py
@@ -4,7 +4,7 @@
 from fastapi.responses import JSONResponse
 from pepdbagent import PEPDatabaseAgent
 
-from sentence_transformers import SentenceTransformer
+from fastembed.embedding import FlagEmbedding as Embedding
 from qdrant_client import QdrantClient
 
 from ....dependencies import (
@@ -29,7 +29,7 @@
 async def search_for_pep(
     query: SearchQuery,
     qdrant: QdrantClient = Depends(get_qdrant),
-    model: SentenceTransformer = Depends(get_sentence_transformer),
+    model: Embedding = Depends(get_sentence_transformer),
     agent: PEPDatabaseAgent = Depends(get_db),
     namespace_access: List[str] = Depends(get_namespace_access_list),
 ):
@@ -42,7 +42,7 @@ async def search_for_pep(
     score_threshold = query.score_threshold
     if qdrant is not None:
         try:
-            query_vec = model.encode(query.query)
+            query_vec = list(model.embed(query.query))[0]
             # count = len(qdrant.search(
             #     collection_name=(
             #         query.collection_name or DEFAULT_QDRANT_COLLECTION_NAME