use BAAI/bge-base-en-v1.5

app.py

@@ -5,8 +5,8 @@ import torch
 from transformers import AutoModel, AutoTokenizer
 import meilisearch
 
-tokenizer = AutoTokenizer.from_pretrained('Snowflake/snowflake-arctic-embed-m')
-model = AutoModel.from_pretrained('Snowflake/snowflake-arctic-embed-m', add_pooling_layer=False)
+tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-base-en-v1.5')
+model = AutoModel.from_pretrained('BAAI/bge-base-en-v1.5')
 model.eval()
 
 cuda_available = torch.cuda.is_available()
@@ -23,16 +23,17 @@ def search_embeddings(query_text):
     # step1: tokenizer the query
     with torch.no_grad():
         # Compute token embeddings
-        query_embeddings = model(**query_tokens)[0][:, 0]
+        model_output = model(**query_tokens)
+        sentence_embeddings = model_output[0][:, 0]
         # normalize embeddings
-        query_embeddings = torch.nn.functional.normalize(query_embeddings, p=2, dim=1)
-        document_embeddings_list = query_embeddings[0].tolist()
+        sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
+        sentence_embeddings_list = sentence_embeddings[0].tolist()
         elapsed_time_embedding = time.time() - start_time_embedding
     
     # step2: search meilisearch
     start_time_meilisearch = time.time()
     response = meilisearch_index.search(
-        "", opt_params={"vector": document_embeddings_list, "hybrid": {"semanticRatio": 1.0}, "limit": 5, "attributesToRetrieve": ["text", "source", "library"]}
+        "", opt_params={"vector": sentence_embeddings_list, "hybrid": {"semanticRatio": 1.0}, "limit": 5, "attributesToRetrieve": ["text", "source", "library"]}
     )
     elapsed_time_meilisearch = time.time() - start_time_meilisearch
     hits = response["hits"]