updates for making similarity work

meedan · Sep 16, 2024 · 0c165d8 · 0c165d8
1 parent d844810
commit 0c165d8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/app/main/lib/elastic_crud.py b/app/main/lib/elastic_crud.py
@@ -42,7 +42,7 @@ def get_presto_request_response(modality, callback_url, task):
 
 def requires_encoding(obj):
     for model_key in obj.get("models", []):
-        if not obj.get('model_'+model_key):
+        if model_key != "elasticsearch" and not obj.get('model_'+model_key):
             return True
     return False
 

diff --git a/app/main/lib/text_similarity.py b/app/main/lib/text_similarity.py
@@ -36,7 +36,7 @@ def async_search_text(task, modality):
     return elastic_crud.get_async_presto_response(task, "text", modality)
 
 def fill_in_openai_embeddings(document):
-    for model_key in document.pop("models", []):
+    for model_key in document.get("models", []):
         if model_key != "elasticsearch" and model_key[:len(PREFIX_OPENAI)] == PREFIX_OPENAI:
             document['vector_'+model_key] = retrieve_openai_embeddings(document['content'], model_key)
             document['model_'+model_key] = 1
@@ -76,7 +76,7 @@ def search_text(search_params, use_document_vectors=False):
     if model_key != "elasticsearch":
       search_params.pop("model", None)
       if use_document_vectors:
-          vector_for_search = search_params[model_key+"-tokens"]
+          vector_for_search = search_params["vector_"+model_key]
       else:
           vector_for_search = None
     result = search_text_by_model(dict(**search_params, **{'model': model_key}), vector_for_search)