import refactoring

Signed-off-by: Daniel Korzekwa <[email protected]>

Author: danielkorzekwa

modelopt/torch/_compress/subblock_stats/calc_subblock_memory.py

@@ -19,8 +19,12 @@
 
 import numpy as np
 import torch
-from puzzle_tools.deci_lm_hf_code.block_config import AttentionConfig, FFNConfig, MambaConfig
 
+from modelopt.torch._compress.decilm.deci_lm_hf_code.block_config import (
+    AttentionConfig,
+    FFNConfig,
+    MambaConfig,
+)
 from modelopt.torch._compress.decilm.deci_lm_hf_code.configuration_decilm import DeciLMConfig
 from modelopt.torch._compress.decilm.deci_lm_hf_code.modeling_decilm import DeciLMMoe
 from modelopt.torch._compress.utils.utils import (

modelopt/torch/_compress/subblock_stats/calc_subblock_stats.py

@@ -34,16 +34,15 @@
 import torch
 from frozendict import frozendict
 from omegaconf import DictConfig, ListConfig, OmegaConf
-from puzzle_tools.deci_lm_hf_code.block_config import (
+from tqdm import tqdm
+
+from modelopt.torch._compress.decilm.deci_lm_hf_code.block_config import (
     AttentionConfig,
     BlockConfig,
     FFNConfig,
     SubblockConfig,
 )
-from puzzle_tools.replacement_utils import parse_layer_replacement
-from tqdm import tqdm
-from utils.parsing import format_global_config
-
+from modelopt.torch._compress.replacement_library.replacement_utils import parse_layer_replacement
 from modelopt.torch._compress.subblock_stats.calc_subblock_memory import (
     calc_subblock_active_params,
     calculate_non_block_memory,
@@ -55,6 +54,7 @@
 from modelopt.torch._compress.tools.hydra_utils import register_hydra_resolvers
 from modelopt.torch._compress.tools.logger import mprint
 from modelopt.torch._compress.tools.robust_json import json_dump
+from modelopt.torch._compress.utils.parsing import format_global_config
 
 # Type variable for dataclasses
 T_DataClass = TypeVar("T_DataClass")