mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-10-02 07:40:28 +00:00
929e7fe96b
1.delete CFG in embedding_engine api 2.add a text_splitter param in embedding_engine api 3.fmt
38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
from typing import Dict, List, Optional
|
|
|
|
from langchain.document_loaders import CSVLoader
|
|
from langchain.schema import Document
|
|
from langchain.text_splitter import TextSplitter
|
|
|
|
from pilot.embedding_engine import SourceEmbedding, register
|
|
|
|
|
|
class CSVEmbedding(SourceEmbedding):
|
|
"""csv embedding for read csv document."""
|
|
|
|
def __init__(
|
|
self,
|
|
file_path,
|
|
vector_store_config,
|
|
text_splitter: Optional[TextSplitter] = None,
|
|
):
|
|
"""Initialize with csv path."""
|
|
super().__init__(file_path, vector_store_config, text_splitter=None)
|
|
self.file_path = file_path
|
|
self.vector_store_config = vector_store_config
|
|
self.text_splitter = text_splitter or None
|
|
|
|
@register
|
|
def read(self):
|
|
"""Load from csv path."""
|
|
loader = CSVLoader(file_path=self.file_path)
|
|
return loader.load()
|
|
|
|
@register
|
|
def data_process(self, documents: List[Document]):
|
|
i = 0
|
|
for d in documents:
|
|
documents[i].page_content = d.page_content.replace("\n", "")
|
|
i += 1
|
|
return documents
|