logo
Browse Source

Add eqa inseert

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
shiyu22 2 years ago
parent
commit
93ce68aea0
  1. 6
      eqa_insert.py

6
eqa_insert.py

@ -26,7 +26,6 @@ class EnhancedQAInsertConfig:
self.source_type = 'file' self.source_type = 'file'
# config for sentence_embedding # config for sentence_embedding
self.model = 'all-MiniLM-L6-v2' self.model = 'all-MiniLM-L6-v2'
self.dim = None
self.openai_api_key = None self.openai_api_key = None
self.customize_embedding_op = None self.customize_embedding_op = None
self.normalize_vec = True self.normalize_vec = True
@ -70,9 +69,6 @@ def _get_embedding_op(config):
@AutoPipes.register @AutoPipes.register
def enhanced_qa_insert_pipe(config): def enhanced_qa_insert_pipe(config):
if not config:
config = EnhancedQAInsertConfig()
text_load_op = ops.text_loader() text_load_op = ops.text_loader()
allow_triton, sentence_embedding_op = _get_embedding_op(config) allow_triton, sentence_embedding_op = _get_embedding_op(config)
@ -97,7 +93,7 @@ def enhanced_qa_insert_pipe(config):
) )
if config.normalize_vec: if config.normalize_vec:
p = p.map('vec', 'vec', ops.towhee.np_normalize())
p = p.map('embedding', 'embedding', ops.towhee.np_normalize())
return (p.map(('doc', 'sentence', 'embedding'), 'mr', insert_milvus_op) return (p.map(('doc', 'sentence', 'embedding'), 'mr', insert_milvus_op)
.output('mr') .output('mr')

Loading…
Cancel
Save