github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-09-19 00:14:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

ci: make ci happy lint the code, delete unused imports

Browse Source 
Signed-off-by: yihong0618 <zouzou0208@gmail.com>
This commit is contained in:
yihong0618
2023-05-24 18:42:55 +08:00
parent 562d5a98cc
commit b098a48898
75 changed files with 1110 additions and 824 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
tools/knowlege_init.py
View File

@@ -2,8 +2,13 @@
# -*- 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.model_config import (
DATASETS_DIR,
LLM_MODEL_CONFIG,
VECTOR_SEARCH_TOP_K,
VECTOR_STORE_CONFIG,
VECTOR_STORE_TYPE,
)
from pilot.source_embedding.knowledge_embedding import KnowledgeEmbedding
@@ -16,22 +21,24 @@ class LocalKnowledgeInit:
self.vector_store_config = vector_store_config
def knowledge_persist(self, file_path, append_mode):
""" knowledge persist """
"""knowledge persist"""
kv = KnowledgeEmbedding(
file_path=file_path,
model_name=LLM_MODEL_CONFIG["text2vec"],
vector_store_config= self.vector_store_config)
vector_store_config=self.vector_store_config,
)
vector_store = kv.knowledge_persist_initialization(append_mode)
return vector_store
def query(self, q):
"""Query similar doc from Vector """
"""Query similar doc from Vector"""
vector_store = self.init_vector_store()
docs = vector_store.similarity_search_with_score(q, k=self.top_k)
for doc in docs:
dc, s = doc
yield s, dc
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--vector_name", type=str, default="default")
@@ -41,8 +48,12 @@ if __name__ == "__main__":
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}
vector_store_config = {
"url": VECTOR_STORE_CONFIG["url"],
"port": VECTOR_STORE_CONFIG["port"],
"vector_store_name": vector_name,
}
print(vector_store_config)
kv = LocalKnowledgeInit(vector_store_config=vector_store_config)
kv = LocalKnowledgeInit(vector_store_config=vector_store_config)
vector_store = kv.knowledge_persist(file_path=DATASETS_DIR, append_mode=append_mode)
print("your knowledge embedding success...")
print("your knowledge embedding success...")