parsing

Parsing and formatting utilities for configuration handling in model compression.

This module provides utilities for: - Parsing command-line arguments and configuration strings - Formatting and displaying model configurations (block configs, attention, FFN) - Formatting loss metrics for logging and visualization

Functions

handle_arg_string

simple_parse_args_string

Parse args1=val1,arg2=val2 into a dictionary.

parse_json

parse_path

get_nested_key

If nested_key is "a.b.c" returns dictionary["a"]["b"]["c"]

format_global_config

Pretty prints a global DictConfig with nice formatting and visual indicators.

format_stitched_losses

Pretty prints stitched module losses with comprehensive tracking and visual indicators.
format_global_config(config, title='Global Configuration')

Pretty prints a global DictConfig with nice formatting and visual indicators.

Parameters:

  • config (DictConfig) – DictConfig object to format

  • title (str) – Title to display at the top of the formatted output

Returns:

Formatted string with configuration details

Return type:

str

Example output:

╭─────────────────── Global Configuration ────────────────────╮ │ Training │ │ • learning_rate: 1e-4 │ │ • batch_size: 32 │ │ • epochs: 100 │ │ Model │ │ • hidden_dim: 512 │ │ • num_layers: 6 │ │ Data │ │ • dataset_path: /path/to/data │ │ • block_size: 2048 │ ╰──────────────────────────────────────────────────────────────╯

format_stitched_losses(losses_dict, best_steps_dict=None, best_values_dict=None, step_number=None, title='Stitched Module Losses')

Pretty prints stitched module losses with comprehensive tracking and visual indicators.

Parameters:

  • losses_dict (dict[str, float]) – Dictionary with block names as keys and current loss values as floats

  • best_steps_dict (dict[str, int] | None) – Optional dictionary with block names as keys and best step numbers as values

  • best_values_dict (dict[str, float] | None) – Optional dictionary with block names as keys and best loss values as floats

  • step_number (int | None) – Optional current step number to include in summary

  • title (str) – Title to display at the top of the formatted output

Returns:

Formatted string with loss values in a comprehensive table format

Return type:

str

Example output:

╭─────────────────── Stitched Module Losses ──────────────────╮ │ Block │ Loss Value │ Best Step │ Best Value │ Change from avg │ │───────┼────────────┼───────────┼────────────┼──────────────────│ │ 00 │ 6.21e-03 │ Step 5 │ 5.95e-03 │ ↑ +2.6e-04 │ │ 01 │ 5.14e-04 │ Step 12 │ 5.14e-04 │ ↓ -1.2e-04 │ │ 02 │ 9.84e-05 │ Step 15 │ 9.84e-05 │ ↓ -3.1e-04 │ ╰──────────────────────────────────────────────────────────────╯

get_nested_key(dictionary, nested_key)

If nested_key is “a.b.c” returns dictionary[“a”][“b”][“c”]

Parameters:

  • dictionary (dict[str, Any])

  • nested_key (str)

Return type:

Any

handle_arg_string(arg)
parse_json(s)
Parameters:

s (str | None)

Return type:

Any

parse_path(s)
Parameters:

s (str | None)

Return type:

Path | None

simple_parse_args_string(args_string)

Parse args1=val1,arg2=val2 into a dictionary.