github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-09-03 10:05:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

test:add document embedding and url embedding test

Browse Source
This commit is contained in:
aries_ckt
2023-08-02 15:55:16 +08:00
parent f270283091
commit bc2db7347e
2 changed files with 21 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
tests/unit/embedding_engine/document_test.py Normal file
View File

@@ -0,0 +1,20 @@
from pilot 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()

2
tests/unit/embedding_engine/test_url_embedding.py → tests/unit/embedding_engine/url_test.py
View File

@@ -1,7 +1,7 @@
from pilot import EmbeddingEngine, KnowledgeType from pilot import EmbeddingEngine, KnowledgeType
url = "https://db-gpt.readthedocs.io/en/latest/getting_started/getting_started.html" url = "https://db-gpt.readthedocs.io/en/latest/getting_started/getting_started.html"
embedding_model = "text2vec" embedding_model = "your_embedding_model"
vector_store_type = "Chroma" vector_store_type = "Chroma"
chroma_persist_path = "your_persist_path" chroma_persist_path = "your_persist_path"
vector_store_config = { vector_store_config = {