fix

src/llmcompressor/transformers/finetune/data/cnn_dailymail.py

@@ -5,7 +5,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.arg_parser import DatasetArguments
+    from llmcompressor.args import DatasetArguments
 
 
 @TextGenerationDataset.register(name="cnn_dailymail")

src/llmcompressor/transformers/finetune/data/evolcodealpaca.py

@@ -5,7 +5,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.arg_parser import DatasetArguments
+    from llmcompressor.args import DatasetArguments
 
 
 @TextGenerationDataset.register(name="evolcodealpaca")

src/llmcompressor/transformers/finetune/data/flickr_30k.py

@@ -7,7 +7,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.arg_parser import DatasetArguments
+    from llmcompressor.args import DatasetArguments
 
 
 @TextGenerationDataset.register(name="flickr", alias="flickr30k")

src/llmcompressor/transformers/finetune/data/gsm8k.py

@@ -5,7 +5,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.arg_parser import DatasetArguments
+    from llmcompressor.args import DatasetArguments
 
 
 @TextGenerationDataset.register(name="gsm8k")

src/llmcompressor/transformers/finetune/data/open_platypus.py

@@ -5,7 +5,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.arg_parser import DatasetArguments
+    from llmcompressor.args import DatasetArguments
 
 
 @TextGenerationDataset.register(name="open_platypus")

src/llmcompressor/transformers/finetune/data/ptb.py

@@ -5,7 +5,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.arg_parser import DatasetArguments
+    from llmcompressor.args import DatasetArguments
 
 
 @TextGenerationDataset.register(name="ptb")

src/llmcompressor/transformers/finetune/data/ultrachat_200k.py

@@ -7,7 +7,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.arg_parser import DatasetArguments
+    from llmcompressor.args import DatasetArguments
 
 
 @TextGenerationDataset.register(name="ultrachat_200k")

src/llmcompressor/transformers/finetune/data/wikitext.py

@@ -5,7 +5,7 @@
 from llmcompressor.typing import Processor
 
 if TYPE_CHECKING:
-    from llmcompressor.arg_parser import DatasetArguments
+    from llmcompressor.args import DatasetArguments
 
 
 @TextGenerationDataset.register(name="wikitext")

src/llmcompressor/transformers/tracing/debug.py

@@ -10,7 +10,7 @@
 from llmcompressor.utils.pytorch.module import get_no_split_params
 from llmcompressor.pipelines.sequential.helpers import trace_subgraphs
 from llmcompressor.transformers import TextGenerationDataset
-from llmcompressor.arg_parser import DatasetArguments
+from llmcompressor.args import DatasetArguments
 
 
 def parse_args():

src/llmcompressor/transformers/utils/helpers.py

@@ -14,7 +14,7 @@
 from transformers.trainer_utils import get_last_checkpoint
 
 if TYPE_CHECKING:
-    from llmcompressor.arg_parser import ModelArguments, TrainingArguments
+    from llmcompressor.args import ModelArguments, TrainingArguments
 
 __all__ = [
     "RECIPE_FILE_NAME",

tests/llmcompressor/transformers/finetune/data/conftest.py

@@ -1,7 +1,7 @@
 import pytest
 from transformers import AutoTokenizer
 
-from llmcompressor.arg_parser import ModelArguments
+from llmcompressor.args import ModelArguments
 
 
 @pytest.fixture

tests/llmcompressor/transformers/finetune/data/test_dataset_helpers.py

@@ -1,6 +1,6 @@
 import pytest
 
-from llmcompressor.arg_parser import DatasetArguments
+from llmcompressor.args import DatasetArguments
 from llmcompressor.transformers.finetune.data.data_helpers import (
     get_raw_dataset,
     make_dataset_splits,

tests/llmcompressor/transformers/finetune/data/test_dataset_loading.py

@@ -5,7 +5,7 @@
 from datasets import IterableDataset, load_dataset
 from parameterized import parameterized
 
-from llmcompressor.arg_parser import (
+from llmcompressor.args import (
     DatasetArguments,
     ModelArguments,
     RecipeArguments,

tests/llmcompressor/transformers/finetune/data/test_registry.py

@@ -1,6 +1,6 @@
 import pytest
 
-from llmcompressor.arg_parser import DatasetArguments
+from llmcompressor.args import DatasetArguments
 from llmcompressor.transformers.finetune.data import (
     C4Dataset,
     OpenPlatypusDataset,

tests/llmcompressor/transformers/obcq/test_obcq_completion.py

@@ -22,7 +22,7 @@ def labeled_dataloader(self, dataset_name, model_name):
         from torch.utils.data import DataLoader
         from transformers import AutoTokenizer, DefaultDataCollator
 
-        from llmcompressor.arg_parser import DatasetArguments
+        from llmcompressor.args import DatasetArguments
         from llmcompressor.transformers.finetune.data import TextGenerationDataset
 
         tokenizer = AutoTokenizer.from_pretrained(model_name)