zhichyu commited on
Commit
e23293e
·
1 Parent(s): d3fd174

Detect shape error of embedding (#3710)

Browse files

### What problem does this PR solve?

Detect shape error of embedding. Close #2997

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

Files changed (1) hide show
  1. rag/nlp/search.py +3 -0
rag/nlp/search.py CHANGED
@@ -46,6 +46,9 @@ class Dealer:
46
 
47
  def get_vector(self, txt, emb_mdl, topk=10, similarity=0.1):
48
  qv, _ = emb_mdl.encode_queries(txt)
 
 
 
49
  embedding_data = [float(v) for v in qv]
50
  vector_column_name = f"q_{len(embedding_data)}_vec"
51
  return MatchDenseExpr(vector_column_name, embedding_data, 'float', 'cosine', topk, {"similarity": similarity})
 
46
 
47
  def get_vector(self, txt, emb_mdl, topk=10, similarity=0.1):
48
  qv, _ = emb_mdl.encode_queries(txt)
49
+ shape = np.array(qv).shape
50
+ if len(shape) > 1:
51
+ raise Exception(f"Dealer.get_vector returned array's shape {shape} doesn't match expectation(exact one dimension).")
52
  embedding_data = [float(v) for v in qv]
53
  vector_column_name = f"q_{len(embedding_data)}_vec"
54
  return MatchDenseExpr(vector_column_name, embedding_data, 'float', 'cosine', topk, {"similarity": similarity})