DB-GPT/tests/unit_tests/embedding_engine/document_test.py
FangYin Cheng cd725db1fb
refactor: The first refactored version for sdk release (#907)
Co-authored-by: chengfangyin2 <chengfangyin3@jd.com>
2023-12-08 14:45:59 +08:00

22 lines
683 B
Python

from dbgpt import EmbeddingEngine, KnowledgeType
embedding_model = "your_embedding_model"
vector_store_type = "Chroma"
chroma_persist_path = "your_persist_path"
vector_store_config = {
"vector_store_name": "document_test",
"vector_store_type": vector_store_type,
"chroma_persist_path": chroma_persist_path,
}
# it can be .md,.pdf,.docx, .csv, .html
document_path = "your_path/test.md"
embedding_engine = EmbeddingEngine(
knowledge_source=document_path,
knowledge_type=KnowledgeType.DOCUMENT.value,
model_name=embedding_model,
vector_store_config=vector_store_config,
)
# embedding document content to vector store
embedding_engine.knowledge_embedding()