diff --git a/aria/data.py b/aria/data.py
index 1aa2094..1b7459f 100644
--- a/aria/data.py
+++ b/aria/data.py
@@ -22,9 +22,8 @@
 from typing import Dict, Iterable, List
 
 import torch
-from datasets.features import Features, Sequence, Value
-
 from datasets import DatasetDict, concatenate_datasets, load_dataset
+from datasets.features import Features, Sequence, Value
 
 
 def apply_chat_template_and_tokenize(