style: code format.

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

View File

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

View File

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