From 99c8859eb02d1c4bde465d1fb0290a8024d2fd87 Mon Sep 17 00:00:00 2001
From: NanoCode012 <nano@axolotl.ai>
Date: Tue, 27 May 2025 10:44:01 +0700
Subject: [PATCH] chore: refactor

---
 src/axolotl/cli/args.py | 31 +++++++++++++++----------------
 1 file changed, 15 insertions(+), 16 deletions(-)

diff --git a/src/axolotl/cli/args.py b/src/axolotl/cli/args.py
index 088e337e4..35f7870be 100644
--- a/src/axolotl/cli/args.py
+++ b/src/axolotl/cli/args.py
@@ -1,7 +1,6 @@
 """Module for axolotl CLI command arguments."""
 
 from dataclasses import dataclass, field
-from typing import Optional
 
 
 @dataclass
@@ -11,9 +10,9 @@ class PreprocessCliArgs:
     debug: bool = field(default=False)
     debug_text_only: bool = field(default=False)
     debug_num_examples: int = field(default=1)
-    prompter: Optional[str] = field(default=None)
-    download: Optional[bool] = field(default=True)
-    iterable: Optional[bool] = field(
+    prompter: str | None = field(default=None)
+    download: bool | None = field(default=True)
+    iterable: bool | None = field(
         default=None,
         metadata={
             "help": "Use IterableDataset for streaming processing of large datasets"
@@ -29,29 +28,29 @@ class TrainerCliArgs:
     debug_text_only: bool = field(default=False)
     debug_num_examples: int = field(default=0)
     merge_lora: bool = field(default=False)
-    prompter: Optional[str] = field(default=None)
+    prompter: str | None = field(default=None)
     shard: bool = field(default=False)
-    main_process_port: Optional[int] = field(default=None)
-    num_processes: Optional[int] = field(default=None)
+    main_process_port: int | None = field(default=None)
+    num_processes: int | None = field(default=None)
 
 
 @dataclass
 class VllmServeCliArgs:
     """Dataclass with CLI arguments for `axolotl vllm-serve` command."""
 
-    tensor_parallel_size: Optional[int] = field(
+    tensor_parallel_size: int | None = field(
         default=None,
         metadata={"help": "Number of tensor parallel workers to use."},
     )
-    host: Optional[str] = field(
+    host: str | None = field(
         default=None,  # nosec B104
         metadata={"help": "Host address to run the server on."},
     )
-    port: Optional[int] = field(
+    port: int | None = field(
         default=None,
         metadata={"help": "Port to run the server on."},
     )
-    gpu_memory_utilization: Optional[float] = field(
+    gpu_memory_utilization: float | None = field(
         default=None,
         metadata={
             "help": "Ratio (between 0 and 1) of GPU memory to reserve for the model weights, activations, and KV "
@@ -60,14 +59,14 @@ class VllmServeCliArgs:
             "out-of-memory (OOM) errors during initialization."
         },
     )
-    dtype: Optional[str] = field(
+    dtype: str | None = field(
         default=None,
         metadata={
             "help": "Data type to use for vLLM generation. If set to 'auto', the data type will be automatically "
             "determined based on the model configuration. Find the supported values in the vLLM documentation."
         },
     )
-    max_model_len: Optional[int] = field(
+    max_model_len: int | None = field(
         default=None,
         metadata={
             "help": "If set, the `max_model_len` to use for vLLM. This can be useful when running with reduced "
@@ -75,14 +74,14 @@ class VllmServeCliArgs:
             "context size, which might be much larger than the KV cache, leading to inefficiencies."
         },
     )
-    enable_prefix_caching: Optional[bool] = field(
+    enable_prefix_caching: bool | None = field(
         default=None,
         metadata={
             "help": "Whether to enable prefix caching in vLLM. If set to `True`, ensure that the model and the "
             "hardware support this feature."
         },
     )
-    serve_module: Optional[str] = field(
+    serve_module: str | None = field(
         default=None,
         metadata={
             "help": "Module to serve. If not set, the default module will be used."
@@ -103,4 +102,4 @@ class EvaluateCliArgs:
 class InferenceCliArgs:
     """Dataclass with CLI arguments for `axolotl inference` command."""
 
-    prompter: Optional[str] = field(default=None)
+    prompter: str | None = field(default=None)