feature:url,csv embedding

This commit is contained in:
chenketing
2023-05-11 23:48:56 +08:00
parent ed855df01d
commit d42a9f3bd1
7 changed files with 86 additions and 108 deletions

View File

@@ -0,0 +1,33 @@
from typing import List, Optional, Dict
from pilot.source_embedding import SourceEmbedding, register
from langchain.document_loaders import CSVLoader
from langchain.schema import Document
class CSVEmbedding(SourceEmbedding):
"""csv embedding for read csv document."""
def __init__(self, file_path, model_name, vector_store_config, embedding_args: Optional[Dict] = None):
"""Initialize with csv path."""
self.file_path = file_path
self.model_name = model_name
self.vector_store_config = vector_store_config
self.embedding_args = embedding_args
@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