Browse Source
Add eqa inseert
Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
1 changed files with
1 additions and
5 deletions
-
eqa_insert.py
|
|
|
@ -26,7 +26,6 @@ class EnhancedQAInsertConfig: |
|
|
|
self.source_type = 'file' |
|
|
|
# config for sentence_embedding |
|
|
|
self.model = 'all-MiniLM-L6-v2' |
|
|
|
self.dim = None |
|
|
|
self.openai_api_key = None |
|
|
|
self.customize_embedding_op = None |
|
|
|
self.normalize_vec = True |
|
|
|
@ -70,9 +69,6 @@ def _get_embedding_op(config): |
|
|
|
|
|
|
|
@AutoPipes.register |
|
|
|
def enhanced_qa_insert_pipe(config): |
|
|
|
if not config: |
|
|
|
config = EnhancedQAInsertConfig() |
|
|
|
|
|
|
|
text_load_op = ops.text_loader() |
|
|
|
|
|
|
|
allow_triton, sentence_embedding_op = _get_embedding_op(config) |
|
|
|
@ -97,7 +93,7 @@ def enhanced_qa_insert_pipe(config): |
|
|
|
) |
|
|
|
|
|
|
|
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) |
|
|
|
.output('mr') |
|
|
|
|
