update:config env

This commit is contained in:
aries-ckt
2023-05-24 15:22:10 +08:00
parent e6339b06ad
commit 977a88509e
3 changed files with 12 additions and 18 deletions

View File

@@ -2,11 +2,11 @@
# -*- coding: utf-8 -*-
import argparse
from pilot.configs.model_config import DATASETS_DIR, LLM_MODEL_CONFIG, VECTOR_SEARCH_TOP_K, VECTOR_STORE_CONFIG, \
VECTOR_STORE_TYPE
from pilot.configs.config import Config
from pilot.configs.model_config import DATASETS_DIR, LLM_MODEL_CONFIG, VECTOR_SEARCH_TOP_K
from pilot.source_embedding.knowledge_embedding import KnowledgeEmbedding
CFG = Config()
class LocalKnowledgeInit:
embeddings: object = None
model_name = LLM_MODEL_CONFIG["text2vec"]
@@ -32,6 +32,7 @@ class LocalKnowledgeInit:
dc, s = doc
yield s, dc
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--vector_name", type=str, default="default")
@@ -40,8 +41,8 @@ if __name__ == "__main__":
args = parser.parse_args()
vector_name = args.vector_name
append_mode = args.append
store_type = VECTOR_STORE_TYPE
vector_store_config = {"url": VECTOR_STORE_CONFIG["url"], "port": VECTOR_STORE_CONFIG["port"], "vector_store_name":vector_name}
store_type = CFG.VECTOR_STORE_TYPE
vector_store_config = {"vector_store_name": vector_name}
print(vector_store_config)
kv = LocalKnowledgeInit(vector_store_config=vector_store_config)
vector_store = kv.knowledge_persist(file_path=DATASETS_DIR, append_mode=append_mode)