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

feature:add milvus store

Browse Source
This commit is contained in:
aries-ckt
2023-05-21 16:29:00 +08:00
parent 98d50b1b98
commit 6747d877cc
7 changed files with 277 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
tools/knowlege_init.py
View File

@@ -2,8 +2,10 @@
# -*- coding: utf-8 -*-
import argparse
from pilot.configs.model_config import DATASETS_DIR, LLM_MODEL_CONFIG, VECTOR_SEARCH_TOP_K, \
KNOWLEDGE_UPLOAD_ROOT_PATH
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Milvus
from pilot.configs.model_config import DATASETS_DIR, LLM_MODEL_CONFIG, VECTOR_SEARCH_TOP_K, VECTOR_STORE_CONFIG
from pilot.source_embedding.knowledge_embedding import KnowledgeEmbedding
@@ -12,15 +14,15 @@ class LocalKnowledgeInit:
model_name = LLM_MODEL_CONFIG["text2vec"]
top_k: int = VECTOR_SEARCH_TOP_K
def __init__(self) -> None:
pass
def __init__(self, vector_store_config) -> None:
self.vector_store_config = vector_store_config
def knowledge_persist(self, file_path, vector_name, append_mode):
def knowledge_persist(self, file_path, append_mode):
""" knowledge persist """
kv = KnowledgeEmbedding(
file_path=file_path,
model_name=LLM_MODEL_CONFIG["text2vec"],
vector_store_config= {"vector_store_name":vector_name, "vector_store_path": KNOWLEDGE_UPLOAD_ROOT_PATH})
vector_store_config= self.vector_store_config)
vector_store = kv.knowledge_persist_initialization(append_mode)
return vector_store
@@ -34,11 +36,15 @@ class LocalKnowledgeInit:
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--vector_name", type=str, default="default")
parser.add_argument("--vector_name", type=str, default="keting")
parser.add_argument("--append", type=bool, default=False)
parser.add_argument("--store_type", type=str, default="Chroma")
args = parser.parse_args()
vector_name = args.vector_name
append_mode = args.append
kv = LocalKnowledgeInit()
vector_store = kv.knowledge_persist(file_path=DATASETS_DIR, vector_name=vector_name, append_mode=append_mode)
store_type = args.store_type
vector_store_config = {"url": VECTOR_STORE_CONFIG["url"], "port": VECTOR_STORE_CONFIG["port"], "vector_store_name":vector_name, "vector_store_type":store_type}
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)
print("your knowledge embedding success...")