Skip to content

analyzer

shrinkai.analysis.analyzer

Classes:

Name Description
FeatureAnalyzer

Evaluates how well a student model mimics the teacher's internal hidden states.
FeatureAnalyzerReport

Holds representation analysis results and renders a dashboard.

Classes

FeatureAnalyzer

Evaluates how well a student model mimics the teacher's internal hidden states.

Methods:

Name Description
__init__

Initializes the FeatureAnalyzer.
evaluate

Runs the dataset through both models and computes alignment metrics per layer.
Source code in src/shrinkai/analysis/analyzer.py 
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
class FeatureAnalyzer:
    """Evaluates how well a student model mimics the teacher's internal hidden states."""

    def __init__(
        self,
        teacher_extractor: FeatureExtractor,
        student_extractor: FeatureExtractor,
        device: torch.device | str = "auto",
    ) -> None:
        """Initializes the FeatureAnalyzer.

        Args:
            teacher_extractor: Teacher model wrapped in FeatureExtractor.
            student_extractor: Student model wrapped in FeatureExtractor.
            device: Computing target.
        """
        self.teacher = teacher_extractor
        self.student = student_extractor
        self.device = resolve_device(device)
        self.teacher.to(self.device)
        self.student.to(self.device)

    def evaluate(
        self,
        dataloader: DataLoader,
        metrics: list[str | BaseFeatureMetric] | None = None,
    ) -> FeatureAnalyzerReport:
        """Runs the dataset through both models and computes alignment metrics per layer.

        Args:
            dataloader: Dataloader yielding validation batches.
            metrics: List of metric strings ('cka') or instantiated BaseFeatureMetric objects.

        Returns:
            FeatureAnalyzerReport: Formatted report ready for `.show()`.
        """
        if metrics is None:
            metrics = ["cka"]
        self.teacher.eval()
        self.student.eval()

        active_metrics = []
        for m in metrics:
            if isinstance(m, str):
                if m.lower() not in FEATURE_METRICS:
                    raise ValueError(
                        f"Unknown metric '{m}'. Available: {list(FEATURE_METRICS.keys())}"
                    )
                active_metrics.append(FEATURE_METRICS[m.lower()])
            else:
                active_metrics.append(m)

        accumulated_scores = defaultdict(lambda: defaultdict(float))
        num_batches = 0

        with torch.no_grad():
            for batch_idx, batch in enumerate(dataloader):
                inputs = batch[0].to(self.device)
                num_batches += 1

                _, s_features = self.student(inputs)
                _, t_features = self.teacher(inputs)

                common_keys = set(s_features.keys()).intersection(set(t_features.keys()))

                if batch_idx == 0 and not common_keys:
                    warnings.warn(
                        "No common alias keys between the teacher and the studen."
                        "Check `target_layers` in your FeatureExtractors.",
                        stacklevel=2,
                    )
                    break

                for key in common_keys:
                    for metric in active_metrics:
                        score = metric.compute(s_features[key], t_features[key])
                        accumulated_scores[key][metric.name] += score

        if num_batches == 0:
            raise ValueError("No batches in the provided dataloader.")

        final_scores = defaultdict(dict)
        for layer, metric_dict in accumulated_scores.items():
            for metric_name, total_score in metric_dict.items():
                final_scores[layer][metric_name] = total_score / num_batches

        return FeatureAnalyzerReport(dict(final_scores))
Methods:
__init__ 
__init__(
    teacher_extractor: FeatureExtractor,
    student_extractor: FeatureExtractor,
    device: device | str = "auto",
) -> None

Initializes the FeatureAnalyzer.

Parameters:

Name Type Description Default
teacher_extractor FeatureExtractor

Teacher model wrapped in FeatureExtractor.

 required
student_extractor FeatureExtractor

Student model wrapped in FeatureExtractor.

 required
device device | str

Computing target.

 'auto'
Source code in src/shrinkai/analysis/analyzer.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def __init__(
    self,
    teacher_extractor: FeatureExtractor,
    student_extractor: FeatureExtractor,
    device: torch.device | str = "auto",
) -> None:
    """Initializes the FeatureAnalyzer.

    Args:
        teacher_extractor: Teacher model wrapped in FeatureExtractor.
        student_extractor: Student model wrapped in FeatureExtractor.
        device: Computing target.
    """
    self.teacher = teacher_extractor
    self.student = student_extractor
    self.device = resolve_device(device)
    self.teacher.to(self.device)
    self.student.to(self.device)
evaluate 
evaluate(
    dataloader: DataLoader,
    metrics: list[str | BaseFeatureMetric] | None = None,
) -> FeatureAnalyzerReport

Runs the dataset through both models and computes alignment metrics per layer.

Parameters:

Name Type Description Default
dataloader DataLoader

Dataloader yielding validation batches.

 required
metrics list[str | BaseFeatureMetric] | None

List of metric strings ('cka') or instantiated BaseFeatureMetric objects.

 None

Returns:

Name Type Description
FeatureAnalyzerReport FeatureAnalyzerReport

Formatted report ready for .show().
Source code in src/shrinkai/analysis/analyzer.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def evaluate(
    self,
    dataloader: DataLoader,
    metrics: list[str | BaseFeatureMetric] | None = None,
) -> FeatureAnalyzerReport:
    """Runs the dataset through both models and computes alignment metrics per layer.

    Args:
        dataloader: Dataloader yielding validation batches.
        metrics: List of metric strings ('cka') or instantiated BaseFeatureMetric objects.

    Returns:
        FeatureAnalyzerReport: Formatted report ready for `.show()`.
    """
    if metrics is None:
        metrics = ["cka"]
    self.teacher.eval()
    self.student.eval()

    active_metrics = []
    for m in metrics:
        if isinstance(m, str):
            if m.lower() not in FEATURE_METRICS:
                raise ValueError(
                    f"Unknown metric '{m}'. Available: {list(FEATURE_METRICS.keys())}"
                )
            active_metrics.append(FEATURE_METRICS[m.lower()])
        else:
            active_metrics.append(m)

    accumulated_scores = defaultdict(lambda: defaultdict(float))
    num_batches = 0

    with torch.no_grad():
        for batch_idx, batch in enumerate(dataloader):
            inputs = batch[0].to(self.device)
            num_batches += 1

            _, s_features = self.student(inputs)
            _, t_features = self.teacher(inputs)

            common_keys = set(s_features.keys()).intersection(set(t_features.keys()))

            if batch_idx == 0 and not common_keys:
                warnings.warn(
                    "No common alias keys between the teacher and the studen."
                    "Check `target_layers` in your FeatureExtractors.",
                    stacklevel=2,
                )
                break

            for key in common_keys:
                for metric in active_metrics:
                    score = metric.compute(s_features[key], t_features[key])
                    accumulated_scores[key][metric.name] += score

    if num_batches == 0:
        raise ValueError("No batches in the provided dataloader.")

    final_scores = defaultdict(dict)
    for layer, metric_dict in accumulated_scores.items():
        for metric_name, total_score in metric_dict.items():
            final_scores[layer][metric_name] = total_score / num_batches

    return FeatureAnalyzerReport(dict(final_scores))

FeatureAnalyzerReport

Holds representation analysis results and renders a dashboard.

Methods:

Name Description
__init__

Initializes the report.
show

Renders an interactive formatted analysis table in the console.
Source code in src/shrinkai/analysis/analyzer.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
class FeatureAnalyzerReport:
    """Holds representation analysis results and renders a dashboard."""

    def __init__(self, layer_scores: dict[str, dict[str, float]]) -> None:
        """Initializes the report.

        Args:
            layer_scores: Nested dict `{ "stage_1": {"CKA (Linear)": 0.85}, ... }`.
        """
        self.layer_scores = layer_scores

    def show(self) -> None:
        """Renders an interactive formatted analysis table in the console."""
        if not self.layer_scores:
            print("No data to report.")
            return

        console = Console()
        table = Table(
            title="Representation Alignment Report (Feature Distillation)",
            header_style="bold magenta",
        )

        table.add_column("Layer / Stage Alias", style="bold")

        first_layer = next(iter(self.layer_scores.values()))
        metric_names = list(first_layer.keys())
        for metric in metric_names:
            table.add_column(metric, justify="right")

        for layer_name, metrics in self.layer_scores.items():
            row_data = [layer_name]
            for metric in metric_names:
                score = metrics.get(metric, 0.0)
                color = "green" if score >= 0.8 else "yellow" if score >= 0.5 else "red"
                row_data.append(f"[{color}]{score:.3f}[/{color}]")
            table.add_row(*row_data)

        console.print(table)
Methods:
__init__ 
__init__(layer_scores: dict[str, dict[str, float]]) -> None

Initializes the report.

Parameters:

Name Type Description Default
layer_scores dict[str, dict[str, float]]

Nested dict { "stage_1": {"CKA (Linear)": 0.85}, ... }.

 required
Source code in src/shrinkai/analysis/analyzer.py 
17
18
19
20
21
22
23
def __init__(self, layer_scores: dict[str, dict[str, float]]) -> None:
    """Initializes the report.

    Args:
        layer_scores: Nested dict `{ "stage_1": {"CKA (Linear)": 0.85}, ... }`.
    """
    self.layer_scores = layer_scores
show 
show() -> None

Renders an interactive formatted analysis table in the console.

Source code in src/shrinkai/analysis/analyzer.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def show(self) -> None:
    """Renders an interactive formatted analysis table in the console."""
    if not self.layer_scores:
        print("No data to report.")
        return

    console = Console()
    table = Table(
        title="Representation Alignment Report (Feature Distillation)",
        header_style="bold magenta",
    )

    table.add_column("Layer / Stage Alias", style="bold")

    first_layer = next(iter(self.layer_scores.values()))
    metric_names = list(first_layer.keys())
    for metric in metric_names:
        table.add_column(metric, justify="right")

    for layer_name, metrics in self.layer_scores.items():
        row_data = [layer_name]
        for metric in metric_names:
            score = metrics.get(metric, 0.0)
            color = "green" if score >= 0.8 else "yellow" if score >= 0.5 else "red"
            row_data.append(f"[{color}]{score:.3f}[/{color}]")
        table.add_row(*row_data)

    console.print(table)

Functions: