gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Try old version of classify_documents

Browse source
This commit is contained in:
Leon Luithlen 2024-11-12 16:28:22 +01:00
parent 2d74590ec4
commit 83995fa548
1 changed files with 4 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
cognee/tasks/documents/classify_documents.py
View file

@ -10,8 +10,10 @@ EXTENSION_TO_DOCUMENT_CLASS = {
def classify_documents(data_documents: list[Data]) -> list[Document]: def classify_documents(data_documents: list[Data]) -> list[Document]:
documents = [ documents = [
EXTENSION_TO_DOCUMENT_CLASS[data_item.extension](id = data_item.id, title=f"{data_item.name}.{data_item.extension}", raw_data_location=data_item.raw_data_location) PdfDocument(id = data_item.id, name=f"{data_item.name}.{data_item.extension}", raw_data_location=data_item.raw_data_location) if data_item.extension == "pdf" else
AudioDocument(id = data_item.id, name=f"{data_item.name}.{data_item.extension}", raw_data_location=data_item.raw_data_location) if data_item.extension == "audio" else
ImageDocument(id = data_item.id, name=f"{data_item.name}.{data_item.extension}", raw_data_location=data_item.raw_data_location) if data_item.extension == "image" else
TextDocument(id = data_item.id, name=f"{data_item.name}.{data_item.extension}", raw_data_location=data_item.raw_data_location)
for data_item in data_documents for data_item in data_documents
] ]
return documents return documents