diff --git a/cognee/tasks/documents/classify_documents.py b/cognee/tasks/documents/classify_documents.py
index 9ab0726ab..9e3420ab8 100644
--- a/cognee/tasks/documents/classify_documents.py
+++ b/cognee/tasks/documents/classify_documents.py
@@ -1,13 +1,6 @@
 from cognee.modules.data.models import Data
 from cognee.modules.data.processing.document_types import Document, PdfDocument, AudioDocument, ImageDocument, TextDocument
 
-EXTENSION_TO_DOCUMENT_CLASS = {
-    "pdf": PdfDocument,
-    "audio": AudioDocument,
-    "image": ImageDocument,
-    "txt": TextDocument
-}
-
 def classify_documents(data_documents: list[Data]) -> list[Document]:
     documents = [
         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