Refactor CoNLLULoaderComponent to use CoNLLULoader from langchain_community

This commit is contained in:
anovazzi1 2024-01-15 20:08:43 -03:00
commit 26a9d5d670

View file

@ -1,6 +1,6 @@
from langflow import CustomComponent from langflow import CustomComponent
from langchain.docstore.document import Document from langchain.docstore.document import Document
from langflow.field_typing import TemplateField from langchain_community.document_loaders.conllu import CoNLLULoader
class CoNLLULoaderComponent(CustomComponent): class CoNLLULoaderComponent(CustomComponent):
@ -17,15 +17,19 @@ class CoNLLULoaderComponent(CustomComponent):
"file_types": ["conllu"], "file_types": ["conllu"],
"field_type": "file", "field_type": "file",
}, },
"metadata": TemplateField( "metadata": {
display_name="Metadata", "display_name": "Metadata",
required=False, "field_type": "dict",
type="dict", "required": False,
), },
} }
def build(self, file_path: str, metadata: dict) -> Document: def build(self, file_path: str, metadata: dict) -> Document:
# Here, you would use the actual class that loads CoNLL-U files. documents = CoNLLULoader(file_path=file_path).load()
# As I don't have the specific class, I'm returning an instance of Document. if(metadata):
# In a real scenario, you should replace the below Document with the actual loader class. for document in documents:
return Document(file_path=file_path, metadata=metadata) if not document.metadata:
document.metadata = metadata
else:
document.metadata.update(metadata)
return documents