From 43e367f2eabe8082fb0a7aaf90d4222ec6a49e1b Mon Sep 17 00:00:00 2001 From: Zhichang Yu Date: Thu, 28 Nov 2024 14:10:22 +0800 Subject: [PATCH] Detect shape error of embedding (#3710) ### 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) --- rag/nlp/search.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/rag/nlp/search.py b/rag/nlp/search.py index 154d5850d..a24f7a71d 100644 --- a/rag/nlp/search.py +++ b/rag/nlp/search.py @@ -46,6 +46,9 @@ class Dealer: def get_vector(self, txt, emb_mdl, topk=10, similarity=0.1): qv, _ = emb_mdl.encode_queries(txt) + shape = np.array(qv).shape + if len(shape) > 1: + raise Exception(f"Dealer.get_vector returned array's shape {shape} doesn't match expectation(exact one dimension).") embedding_data = [float(v) for v in qv] vector_column_name = f"q_{len(embedding_data)}_vec" return MatchDenseExpr(vector_column_name, embedding_data, 'float', 'cosine', topk, {"similarity": similarity})