logo
Browse Source

Update loader

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
shiyu22 2 years ago
parent
commit
0f53700cee
  1. 53
      loader.py
  2. 1
      requirements.txt

53
loader.py

@ -1,16 +1,49 @@
from towhee.operator import PyOperator
from typing import List, Optional
from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
class TextLoader(PyOperator):
"""
DefaultOperator for _input and _output nodes.
"""
'''Load data from urls or files (paths or file-like objects) as a list of doc chunks'''
def __init__(self,
chunk_size: int = 300,
source_type: str = 'file',
):
self.splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size)
self.source_type = source_type
def __call__(self, data_src) -> List[str]:
if not isinstance(data_src, list):
data_src = [data_src]
if self.source_type == 'file':
docs = self._from_files(data_src)
elif self.source_type == 'url':
docs= self._from_urls(data_src)
else:
raise AttributeError('Invalid source type. Only support "file" or "url".')
docs = self.splitter.split_documents(docs)
return [str(doc.page_content) for doc in docs]
def __init__(self):
#pylint: disable=useless-super-delegation
super().__init__()
def _from_files(self, files: list, encoding: Optional[str] = None) -> List[Document]:
'''Load documents from path or file-like object, return a list of unsplit Langchain Documents'''
docs = []
for file in files:
if hasattr(file, 'name'):
file_path = file.name
else:
file_path = file
with open(file_path, encoding=encoding) as f:
text = f.read()
metadata = {"source": file_path}
docs.append(Document(page_content=text, metadata=metadata))
return docs
def _from_urls(self, urls: List[str]) -> List[Document]:
from langchain.document_loaders import UnstructuredURLLoader
def __call__(self, *args):
if len(args) == 1:
return args[0]
return args
loader = UnstructuredURLLoader(urls=urls)
docs = loader.load()
return docs

1
requirements.txt

@ -0,0 +1 @@
langchain
Loading…
Cancel
Save