File size: 7,113 Bytes
dbc6675
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
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
53
54
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
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
"""

Experiment result tracking and comparison table generation.



Logs results from all tokenizer × domain × model × mode × split combinations

and generates CSV / LaTeX tables for the paper.

"""

import csv
import json
import os
from collections import defaultdict


class ExperimentTracker:
    """

    Track and store experimental results across all configurations.



    Results are stored in a nested dict:

        results[domain][tokenizer][model][mode][split] = metrics_dict

    """

    def __init__(self, output_dir: str = "results"):
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)
        self.results: dict = defaultdict(
            lambda: defaultdict(
                lambda: defaultdict(lambda: defaultdict(lambda: defaultdict(dict)))
            )
        )

    def log_result(

        self,

        domain: str,

        tokenizer: str,

        model: str,

        mode: str,

        split: str,

        metrics: dict,

    ) -> None:
        """

        Log a single experimental result.



        Args:

            domain: e.g. 'blocks', 'gripper'

            tokenizer: e.g. 'wl', 'simhash', 'shortest_path', 'graphbpe', 'random'
            model: e.g. 'lstm', 'xgboost'

            mode: 'state' or 'delta'

            split: 'validation', 'test-interpolation', 'test-extrapolation'

            metrics: dict with keys like 'solved_rate', 'exec_rate', 'vocab_size'

        """
        self.results[domain][tokenizer][model][mode][split] = metrics
        print(
            f"  Logged: {domain}/{tokenizer}/{model}/{mode}/{split} → "
            f"{metrics.get('solved_rate', 'N/A')}"
        )

    def generate_comparison_table(self) -> str:
        """

        Generate a comparison table across all methods.



        Returns: CSV string

        """
        rows = []
        header = [
            "Domain",
            "Tokenizer",
            "Model",
            "Mode",
            "Val Solved",
            "Interp Solved",
            "Extrap Solved",
            "Val Exec",
            "Interp Exec",
            "Extrap Exec",
            "Vocab Size",
        ]
        rows.append(header)

        for domain in sorted(self.results.keys()):
            for tokenizer in sorted(self.results[domain].keys()):
                for model in sorted(self.results[domain][tokenizer].keys()):
                    for mode in sorted(
                        self.results[domain][tokenizer][model].keys()
                    ):
                        splits = self.results[domain][tokenizer][model][mode]
                        row = [domain, tokenizer, model, mode]

                        for split in [
                            "validation",
                            "test-interpolation",
                            "test-extrapolation",
                        ]:
                            m = splits.get(split, {})
                            row.append(f"{m.get('solved_rate', 0):.2%}")

                        for split in [
                            "validation",
                            "test-interpolation",
                            "test-extrapolation",
                        ]:
                            m = splits.get(split, {})
                            row.append(f"{m.get('exec_rate', 0):.2%}")

                        # Vocab size (same for all splits)
                        any_m = next(iter(splits.values()), {})
                        row.append(str(any_m.get("vocab_size", "N/A")))

                        rows.append(row)

        # Write CSV
        csv_path = os.path.join(self.output_dir, "tokenization_comparison.csv")
        with open(csv_path, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerows(rows)
        print(f"Saved comparison table to {csv_path}")

        # Generate LaTeX
        self._generate_latex_table(rows)

        return csv_path

    def _generate_latex_table(self, rows: list[list]) -> None:
        """Generate LaTeX table from rows."""
        header = rows[0]
        data = rows[1:]

        n_cols = len(header)
        col_spec = "l" * n_cols

        lines = [
            r"\begin{table}[ht]",
            r"\centering",
            r"\caption{Comparison of tokenization strategies}",
            r"\label{tab:tokenization_comparison}",
            f"\\begin{{tabular}}{{{col_spec}}}",
            r"\toprule",
            " & ".join(header) + r" \\",
            r"\midrule",
        ]

        for row in data:
            lines.append(" & ".join(str(x) for x in row) + r" \\")

        lines.extend(
            [
                r"\bottomrule",
                r"\end{tabular}",
                r"\end{table}",
            ]
        )

        tex_path = os.path.join(self.output_dir, "tokenization_comparison.tex")
        with open(tex_path, "w") as f:
            f.write("\n".join(lines))
        print(f"Saved LaTeX table to {tex_path}")

    def save_results(self, filepath: str | None = None) -> None:
        """Save all results to JSON."""
        if filepath is None:
            filepath = os.path.join(self.output_dir, "full_results.json")

        # Convert defaultdicts to regular dicts for JSON serialization
        def to_dict(d):
            if isinstance(d, defaultdict):
                return {k: to_dict(v) for k, v in d.items()}
            return d

        with open(filepath, "w") as f:
            json.dump(to_dict(self.results), f, indent=2)
        print(f"Saved full results to {filepath}")

    def load_results(self, filepath: str) -> None:
        """Load results from JSON."""
        with open(filepath, "r") as f:
            data = json.load(f)

        # Populate results
        for domain, tok_data in data.items():
            for tok, model_data in tok_data.items():
                for model, mode_data in model_data.items():
                    for mode, split_data in mode_data.items():
                        for split, metrics in split_data.items():
                            self.results[domain][tok][model][mode][split] = metrics

    def generate_vocabulary_stats(self) -> None:
        """Generate vocabulary size stats per tokenizer per domain."""
        stats = {}
        for domain in self.results:
            for tok in self.results[domain]:
                for model in self.results[domain][tok]:
                    for mode in self.results[domain][tok][model]:
                        for split, m in self.results[domain][tok][model][
                            mode
                        ].items():
                            if "vocab_size" in m:
                                stats.setdefault(tok, {})[domain] = m["vocab_size"]

        stats_path = os.path.join(self.output_dir, "vocabulary_stats.json")
        with open(stats_path, "w") as f:
            json.dump(stats, f, indent=2)
        print(f"Saved vocabulary stats to {stats_path}")