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

style: code format.

Browse Source 
code style format
This commit is contained in:
aries-ckt 2023-06-19 16:58:24 +08:00
parent 3af16cf792
commit c32f3f1766
2 changed files with 19 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
pilot/configs/config.py
View File

@ -152,7 +152,6 @@ class Config(metaclass=Singleton):
self.WEAVIATE_URL = os.getenv("WEAVIATE_URL", "http://127.0.0.1:8080")
# QLoRA
self.QLoRA = os.getenv("QUANTIZE_QLORA", "True")

32
pilot/vector_store/weaviate_store.py
View File

@ -45,18 +45,21 @@ class WeaviateStore(VectorStoreBase):
# }
# vector = self.embedding.embed_query(text)
response = (
self.vector_store_client.query.get(self.vector_name, ["metadata", "page_content"])
# .with_near_vector({"vector": vector})
.with_limit(topk)
.do()
self.vector_store_client.query.get(
self.vector_name, ["metadata", "page_content"]
)
res = response['data']['Get'][list(response['data']['Get'].keys())[0]]
# .with_near_vector({"vector": vector})
.with_limit(topk).do()
)
res = response["data"]["Get"][list(response["data"]["Get"].keys())[0]]
docs = []
for r in res:
docs.append(Document(
page_content=r['page_content'],
metadata={"metadata": r['metadata']},
))
docs.append(
Document(
page_content=r["page_content"],
metadata={"metadata": r["metadata"]},
)
)
return docs
def vector_name_exists(self) -> bool:
@ -110,7 +113,7 @@ class WeaviateStore(VectorStoreBase):
# },
"description": "Text content of the document",
"name": "page_content",
}
},
],
# "vectorizer": "text2vec-transformers",
}
@ -132,7 +135,12 @@ class WeaviateStore(VectorStoreBase):
# Batch import all documents
for i in range(len(texts)):
properties = {"metadata": metadatas[i]['source'], "page_content": texts[i]}
properties = {
"metadata": metadatas[i]["source"],
"page_content": texts[i],
}
self.vector_store_client.batch.add_data_object(data_object=properties, class_name=self.vector_name)
self.vector_store_client.batch.add_data_object(
data_object=properties, class_name=self.vector_name
)
self.vector_store_client.batch.flush()