github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-07-30 15:21:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat:embedding_engine add text_splitter param

Browse Source
This commit is contained in:
aries_ckt 2023-07-12 18:01:22 +08:00
parent 30adbaf4fd
commit 56c1947eda
6 changed files with 53 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
docs/getting_started/installation.md
View File

@ -19,6 +19,7 @@ you will prepare embedding models from huggingface
Notice make sure you have install git-lfs
```{tip}
git clone https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2
git clone https://huggingface.co/GanymedeNil/text2vec-large-chinese
```
version:

18
docs/modules/knowledge.rst
View File

@ -72,6 +72,24 @@ eg: git clone https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2
vector_store_config=vector_store_config)
embedding_engine.knowledge_embedding()
If you want to add your text_splitter, do this:
::
url = "https://db-gpt.readthedocs.io/en/latest/getting_started/getting_started.html"
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=100, chunk_overlap=50
)
embedding_engine = EmbeddingEngine(
knowledge_source=url,
knowledge_type=KnowledgeType.URL.value,
model_name=embedding_model,
vector_store_config=vector_store_config,
text_splitter=text_splitter
)
4.init Document Type EmbeddingEngine api and embedding your document into vector store in your code.
Document type can be .txt, .pdf, .md, .doc, .ppt.

49
docs/use_cases/knownledge_based_qa.md
View File

@ -1,49 +0,0 @@
# Knownledge based qa
Chat with your own knowledge is a very interesting thing. In the usage scenarios of this chapter, we will introduce how to build your own knowledge base through the knowledge base API. Firstly, building a knowledge store can currently be initialized by executing "python tool/knowledge_init.py" to initialize the content of your own knowledge base, which was introduced in the previous knowledge base module. Of course, you can also call our provided knowledge embedding API to store knowledge.
We currently support many document formats: txt, pdf, md, html, doc, ppt, and url.
```
vector_store_config = {
"vector_store_name": name
}
file_path = "your file path"
embedding_engine = EmbeddingEngine(file_path=file_path, model_name=LLM_MODEL_CONFIG["text2vec"], vector_store_config=vector_store_config)
embedding_engine.knowledge_embedding()
```
Now we currently support vector databases: Chroma (default) and Milvus. You can switch between them by modifying the "VECTOR_STORE_TYPE" field in the .env file.
```
#*******************************************************************#
#** VECTOR STORE SETTINGS **#
#*******************************************************************#
VECTOR_STORE_TYPE=Chroma
#MILVUS_URL=127.0.0.1
#MILVUS_PORT=19530
```
Below is an example of using the knowledge base API to query knowledge:
```
vector_store_config = {
"vector_store_name": your_name,
"vector_store_type": "Chroma",
"chroma_persist_path": "your_persist_dir",
}
integrate
query = "your query"
embedding_model = "your_model_path/all-MiniLM-L6-v2"
embedding_engine = EmbeddingEngine(knowledge_source=url, knowledge_type=KnowledgeType.URL.value, model_name=embedding_model, vector_store_config=vector_store_config)
embedding_engine.similar_search(query, 10)
```

8
pilot/embedding_engine/embedding_engine.py
View File

@ -2,6 +2,7 @@ from typing import Optional
from chromadb.errors import NotEnoughElementsException
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import TextSplitter
from pilot.embedding_engine.knowledge_type import get_knowledge_embedding, KnowledgeType
from pilot.vector_store.connector import VectorStoreConnector
@ -21,6 +22,7 @@ class EmbeddingEngine:
vector_store_config,
knowledge_type: Optional[str] = KnowledgeType.DOCUMENT.value,
knowledge_source: Optional[str] = None,
text_splitter: Optional[TextSplitter] = None,
):
"""Initialize with knowledge embedding client, model_name, vector_store_config, knowledge_type, knowledge_source"""
self.knowledge_source = knowledge_source
@ -29,6 +31,7 @@ class EmbeddingEngine:
self.knowledge_type = knowledge_type
self.embeddings = HuggingFaceEmbeddings(model_name=self.model_name)
self.vector_store_config["embeddings"] = self.embeddings
self.text_splitter = text_splitter
def knowledge_embedding(self):
"""source embedding is chain process.read->text_split->data_process->index_store"""
@ -47,7 +50,10 @@ class EmbeddingEngine:
def init_knowledge_embedding(self):
return get_knowledge_embedding(
self.knowledge_type, self.knowledge_source, self.vector_store_config
self.knowledge_type,
self.knowledge_source,
self.vector_store_config,
self.text_splitter,
)
def similar_search(self, text, topk):

7
pilot/embedding_engine/knowledge_type.py
View File

@ -40,7 +40,9 @@ class KnowledgeType(Enum):
YOUTUBE = "YOUTUBE"
def get_knowledge_embedding(knowledge_type, knowledge_source, vector_store_config):
def get_knowledge_embedding(
knowledge_type, knowledge_source, vector_store_config, text_splitter
):
match knowledge_type:
case KnowledgeType.DOCUMENT.value:
extension = "." + knowledge_source.rsplit(".", 1)[-1]
@ -49,6 +51,7 @@ def get_knowledge_embedding(knowledge_type, knowledge_source, vector_store_confi
embedding = knowledge_class(
knowledge_source,
vector_store_config=vector_store_config,
text_splitter=text_splitter,
**knowledge_args,
)
return embedding
@ -57,12 +60,14 @@ def get_knowledge_embedding(knowledge_type, knowledge_source, vector_store_confi
embedding = URLEmbedding(
file_path=knowledge_source,
vector_store_config=vector_store_config,
text_splitter=text_splitter,
)
return embedding
case KnowledgeType.TEXT.value:
embedding = StringEmbedding(
file_path=knowledge_source,
vector_store_config=vector_store_config,
text_splitter=text_splitter,
)
return embedding
case KnowledgeType.OSS.value:

21
pilot/server/knowledge/service.py
View File

@ -1,6 +1,8 @@
import threading
from datetime import datetime
from langchain.text_splitter import RecursiveCharacterTextSplitter, SpacyTextSplitter
from pilot.configs.config import Config
from pilot.configs.model_config import LLM_MODEL_CONFIG, KNOWLEDGE_UPLOAD_ROOT_PATH
from pilot.embedding_engine.embedding_engine import EmbeddingEngine
@ -122,6 +124,24 @@ class KnowledgeService:
raise Exception(
f" doc:{doc.doc_name} status is {doc.status}, can not sync"
)
if CFG.LANGUAGE == "en":
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE,
chunk_overlap=20,
length_function=len,
)
else:
try:
text_splitter = SpacyTextSplitter(
pipeline="zh_core_web_sm",
chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE,
chunk_overlap=100,
)
except Exception:
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=CFG.KNOWLEDGE_CHUNK_SIZE, chunk_overlap=50
)
client = EmbeddingEngine(
knowledge_source=doc.content,
knowledge_type=doc.doc_type.upper(),
@ -131,6 +151,7 @@ class KnowledgeService:
"vector_store_type": CFG.VECTOR_STORE_TYPE,
"chroma_persist_path": KNOWLEDGE_UPLOAD_ROOT_PATH,
},
text_splitter=text_splitter,
)
chunk_docs = client.read()
# update document status