From 8b7c2e417c23a00046e44842f20aeaf9894b0fd5 Mon Sep 17 00:00:00 2001 From: imClumsyPanda Date: Wed, 7 Jun 2023 22:14:08 +0800 Subject: [PATCH] add support for csv file --- chains/local_doc_qa.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/chains/local_doc_qa.py b/chains/local_doc_qa.py index 25016f7..c96ede5 100644 --- a/chains/local_doc_qa.py +++ b/chains/local_doc_qa.py @@ -1,6 +1,6 @@ from langchain.embeddings.huggingface import HuggingFaceEmbeddings from langchain.vectorstores import FAISS -from langchain.document_loaders import UnstructuredFileLoader, TextLoader +from langchain.document_loaders import UnstructuredFileLoader, TextLoader, CSVLoader from configs.model_config import * import datetime from textsplitter import ChineseTextSplitter @@ -74,6 +74,9 @@ def load_file(filepath, sentence_size=SENTENCE_SIZE): loader = UnstructuredPaddleImageLoader(filepath, mode="elements") textsplitter = ChineseTextSplitter(pdf=False, sentence_size=sentence_size) docs = loader.load_and_split(text_splitter=textsplitter) + elif filepath.lower().endswith(".csv"): + loader = CSVLoader(filepath) + docs = loader.load() else: loader = UnstructuredFileLoader(filepath, mode="elements") textsplitter = ChineseTextSplitter(pdf=False, sentence_size=sentence_size)