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

update:vector store config

Browse Source
This commit is contained in:
aries-ckt
2023-05-23 22:06:07 +08:00
parent 90b793f40e
commit ef64935145
5 changed files with 25 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
pilot/source_embedding/knowledge_embedding.py
View File

@@ -3,6 +3,8 @@ import os
from bs4 import BeautifulSoup
from langchain.document_loaders import TextLoader, markdown
from langchain.embeddings import HuggingFaceEmbeddings
from pilot.configs.config import Config
from pilot.configs.model_config import DATASETS_DIR, KNOWLEDGE_CHUNK_SPLIT_SIZE, VECTOR_STORE_TYPE
from pilot.source_embedding.chn_document_splitter import CHNDocumentSplitter
from pilot.source_embedding.csv_embedding import CSVEmbedding
@@ -13,6 +15,7 @@ import markdown
from pilot.source_embedding.pdf_loader import UnstructuredPaddlePDFLoader
from pilot.vector_store.connector import VectorStoreConnector
CFG = Config()
class KnowledgeEmbedding:
def __init__(self, file_path, model_name, vector_store_config, local_persist=True):
@@ -53,7 +56,7 @@ class KnowledgeEmbedding:
def knowledge_persist_initialization(self, append_mode):
documents = self._load_knownlege(self.file_path)
self.vector_client = VectorStoreConnector(VECTOR_STORE_TYPE, self.vector_store_config)
self.vector_client = VectorStoreConnector(CFG.VECTOR_STORE_TYPE, self.vector_store_config)
self.vector_client.load_document(documents)
return self.vector_client

6
pilot/source_embedding/source_embedding.py
View File

@@ -4,10 +4,12 @@ from abc import ABC, abstractmethod
from langchain.embeddings import HuggingFaceEmbeddings
from typing import List, Optional, Dict
from pilot.configs.model_config import VECTOR_STORE_TYPE
from pilot.configs.config import Config
from pilot.vector_store.connector import VectorStoreConnector
registered_methods = []
CFG = Config()
def register(method):
@@ -30,7 +32,7 @@ class SourceEmbedding(ABC):
self.embeddings = HuggingFaceEmbeddings(model_name=self.model_name)
vector_store_config["embeddings"] = self.embeddings
self.vector_client = VectorStoreConnector(VECTOR_STORE_TYPE, vector_store_config)
self.vector_client = VectorStoreConnector(CFG.VECTOR_STORE_TYPE, vector_store_config)
@abstractmethod
@register