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

update:config env

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
pilot/vector_store/milvus_store.py
View File

@ -139,29 +139,21 @@ class MilvusStore(VectorStoreBase):
fields.append( fields.append(
FieldSchema(text_field, DataType.VARCHAR, max_length=max_length + 1) FieldSchema(text_field, DataType.VARCHAR, max_length=max_length + 1)
) )
# create the primary key field # primary key field
fields.append( fields.append(
FieldSchema(primary_field, DataType.INT64, is_primary=True, auto_id=True) FieldSchema(primary_field, DataType.INT64, is_primary=True, auto_id=True)
) )
# create the vector field # vector field
fields.append(FieldSchema(vector_field, DataType.FLOAT_VECTOR, dim=dim)) fields.append(FieldSchema(vector_field, DataType.FLOAT_VECTOR, dim=dim))
# Create the schema for the collection # milvus the schema for the collection
schema = CollectionSchema(fields) schema = CollectionSchema(fields)
# Create the collection # Create the collection
collection = Collection(collection_name, schema) collection = Collection(collection_name, schema)
self.col = collection self.col = collection
# Index parameters for the collection # index parameters for the collection
index = self.index_params index = self.index_params
# Create the index # milvus index
collection.create_index(vector_field, index) collection.create_index(vector_field, index)
# Create the VectorStore
# milvus = cls(
# embedding,
# kwargs.get("connection_args", {"port": 19530}),
# collection_name,
# text_field,
# )
# Add the texts.
schema = collection.schema schema = collection.schema
for x in schema.fields: for x in schema.fields:
self.fields.append(x.name) self.fields.append(x.name)

1
requirements.txt
View File

@ -69,6 +69,7 @@ colorama
playsound playsound
distro distro
pypdf pypdf
milvus-cli==0.3.2
# Testing dependencies # Testing dependencies
pytest pytest

11
tools/knowlege_init.py
View File

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