Tom Aarsen commited on
Commit
92d5d6d
·
1 Parent(s): 4d9f5e6

Resolve search crash due to wrong shape

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -28,12 +28,12 @@ model = SentenceTransformer(
28
  def search(query, top_k: int = 10, rerank_multiplier: int = 4):
29
  # 1. Embed the query as float32
30
  start_time = time.time()
31
- query_embedding = model.encode([query])
32
  embed_time = time.time() - start_time
33
 
34
  # 2. Quantize the query to ubinary
35
  start_time = time.time()
36
- query_embedding_ubinary = quantize_embeddings(query_embedding, "ubinary")
37
  quantize_time = time.time() - start_time
38
 
39
  # 3. Search the binary index
 
28
  def search(query, top_k: int = 10, rerank_multiplier: int = 4):
29
  # 1. Embed the query as float32
30
  start_time = time.time()
31
+ query_embedding = model.encode(query)
32
  embed_time = time.time() - start_time
33
 
34
  # 2. Quantize the query to ubinary
35
  start_time = time.time()
36
+ query_embedding_ubinary = quantize_embeddings(query_embedding.reshape(1, -1), "ubinary")
37
  quantize_time = time.time() - start_time
38
 
39
  # 3. Search the binary index