66"""
77
88from dataclasses import dataclass , field
9- from typing import Any , List , Optional
9+ from typing import Any
1010
1111from loguru import logger
1212
@@ -23,11 +23,11 @@ class CompressionLifecycle:
2323 A class for managing the lifecycle of compression events in the LLM Compressor.
2424
2525 :param state: The current state of the compression process
26- :type state: Optional[ State]
26+ :type state: State
2727 :param recipe: The compression recipe
2828 :type recipe: Recipe
2929 :param modifiers: The list of stage modifiers
30- :type modifiers: List [StageModifiers]
30+ :type modifiers: list [StageModifiers]
3131 """
3232
3333 state : State = field (default_factory = State )
@@ -37,8 +37,8 @@ class CompressionLifecycle:
3737 finalized : bool = False
3838
3939 # event order validation
40- _last_event_type : Optional [ EventType ] = EventType .BATCH_END
41- _event_order : List [EventType ] = field (
40+ _last_event_type : EventType | None = EventType .BATCH_END
41+ _event_order : list [EventType ] = field (
4242 default_factory = lambda : [
4343 EventType .BATCH_START ,
4444 EventType .LOSS_CALCULATED ,
@@ -72,11 +72,11 @@ def reset(self):
7272
7373 def initialize (
7474 self ,
75- recipe : Optional [ RecipeInput ] = None ,
76- recipe_stage : Optional [ RecipeStageInput ] = None ,
77- recipe_args : Optional [ RecipeArgsInput ] = None ,
75+ recipe : RecipeInput | None = None ,
76+ recipe_stage : RecipeStageInput | None = None ,
77+ recipe_args : RecipeArgsInput | None = None ,
7878 ** kwargs ,
79- ) -> List [Any ]:
79+ ) -> list [Any ]:
8080 """
8181 Initialize the compression lifecycle.
8282
@@ -114,7 +114,7 @@ def initialize(
114114
115115 return mod_data
116116
117- def finalize (self , ** kwargs ) -> List [Any ]:
117+ def finalize (self , ** kwargs ) -> list [Any ]:
118118 """
119119 Finalize the compression lifecycle.
120120
@@ -149,8 +149,8 @@ def finalize(self, **kwargs) -> List[Any]:
149149 return mod_data
150150
151151 def event (
152- self , event_type : EventType , global_step : Optional [ int ] = 0 , ** kwargs
153- ) -> List [Any ]:
152+ self , event_type : EventType , global_step : int | None = 0 , ** kwargs
153+ ) -> list [Any ]:
154154 """
155155 Handle a compression event.
156156
0 commit comments