From d26bbe19f4b0b694d6bca560b51cd60e42a36cbc Mon Sep 17 00:00:00 2001 From: chaudhryumer Date: Tue, 11 Aug 2026 00:04:20 -0700 Subject: [PATCH 1/2] fix: regenerated splits with OP:partial fix & updated train config --- config.json | 13 +++-- model/transformer.py | 41 +++------------ train.py | 118 +++++++++++++++++++++++++++---------------- 3 files changed, 89 insertions(+), 83 deletions(-) diff --git a/config.json b/config.json index dbfd64d..0b04d6e 100644 --- a/config.json +++ b/config.json @@ -1,12 +1,15 @@ { "learning_rate": 0.0001, + "warmup_steps": 1000, "batch_size": 32, "max_steps": 3500, "hidden_dim": 256, "max_len": 32, - "epochs": 15, - "early_stopping": { - "patience": 12, - "min_delta": 0.0002}, + "epochs": 10, + "grad_clip_max_norm": 1.0, + "early_stopping": { + "patience": 12, + "min_delta": 0.0002 + }, "validation_logging": true -} +} \ No newline at end of file diff --git a/model/transformer.py b/model/transformer.py index b1ff40f..80c4b6f 100644 --- a/model/transformer.py +++ b/model/transformer.py @@ -10,6 +10,7 @@ def __init__( self, vocab_size: int, num_rules: int, + rule_labels: list = None, # Handled rule_labels dynamically hidden_dim: int = 128, num_heads: int = 8, num_layers: int = 8, @@ -19,16 +20,6 @@ def __init__( pad_id: int = 0, ): super().__init__() - # FIX 1 (docs/KNOWN_ISSUES.md, "RuleHead only pools from token 0"): - # RuleHead.forward() falls back to encoder_out[:, 0, :] whenever no - # root_mask is supplied -- i.e. it bases every rule prediction on the - # encoder's representation of ONLY the first source token, ignoring - # the rest of the expression entirely. This was the likely root - # cause of the persistent ~0.505 Val Rule loss plateau seen across - # every training configuration tried (learning rate, data coverage, - # rule-label conflation fix, gradient clipping, hidden_dim increase - # all failed to break it). pad_id is now stored so forward() can - # build a real root_mask covering all non-padding tokens. self.pad_id = pad_id self.encoder = TreeEncoder( @@ -41,8 +32,10 @@ def __init__( position_dim=position_dim, ) - # Instantiate rule labels based on num_rules - rule_labels = [f"RULE_{i}" for i in range(num_rules)] + # Dynamic rule label mapping (resolves RULE_i placeholder issue) + if rule_labels is None: + rule_labels = [f"RULE:{i}" for i in range(num_rules)] + self.rule_head = RuleHead( hidden_dim=hidden_dim, rule_labels=rule_labels @@ -57,9 +50,6 @@ def __init__( dropout=dropout, ) - # In train.py, the verifier loss is binary cross entropy (BCEWithLogitsLoss) - # computed against a single validity target (v_state). Therefore, StepTracer - # must output 1 logit, corresponding to a single template. templates = ["is_valid"] self.step_tracer = StepTracer( hidden_dim=hidden_dim, @@ -70,7 +60,6 @@ def forward(self, src_seq, tgt_in_seq, true_rule_ids=None): device = src_seq.device batch_size, seq_len = src_seq.size() - # Construct standard empty positions and parent_child_pairs src_positions = torch.zeros( (batch_size, seq_len, 3), dtype=torch.float32, device=device ) @@ -83,29 +72,11 @@ def forward(self, src_seq, tgt_in_seq, true_rule_ids=None): src_seq, src_positions, parent_child_pairs ) - # 2. Get rule logits - # FIX 1: build a root_mask covering every real (non-padding) source - # token, instead of letting RuleHead silently fall back to pooling - # only from position 0. This lets the rule classifier actually see - # the whole expression (operator, operand, coefficients, structure) - # rather than just the first token (e.g. "OP:diff"), which is - # identical across many semantically different problems and cannot - # by itself distinguish them. + # 2. Get rule logits (using non-pad tokens root mask) root_mask = (src_seq != self.pad_id) rule_logits = self.rule_head(encoder_output, root_mask=root_mask) # 3. Embed rule IDs for decoder - # FIX 2 (docs/KNOWN_ISSUES.md, "rule/decoder circular dependency"): - # model/architecture.py's older CalculusModel already demonstrates - # this exact pattern -- an optional true_rule_ids parameter that, - # when supplied (training), is used instead of the model's own - # (possibly wrong) argmax prediction. Without this, the decoder was - # always conditioned on the rule head's own guess even during - # training, meaning a wrong early rule prediction corrupted the - # decoder's training signal too, and neither component could - # specialize independently. At inference time (true_rule_ids=None, - # the default), behavior is unchanged -- the model still falls back - # to its own prediction, exactly as before. if true_rule_ids is not None: rule_ids = true_rule_ids else: diff --git a/train.py b/train.py index 407d041..29cac2d 100644 --- a/train.py +++ b/train.py @@ -1,9 +1,11 @@ import sys import os import json +import subprocess import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader +from torch.optim.lr_scheduler import LambdaLR from pathlib import Path sys.path.insert(0, os.path.abspath(os.path.dirname(__file__))) @@ -15,6 +17,15 @@ config = json.load(cfg_file) +def get_git_commit_hash(): + """Returns the exact current git commit hash for provenance tracking.""" + try: + hash_str = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("utf-8").strip() + return hash_str + except Exception: + return "UNKNOWN_COMMIT" + + def flatten_vocab(raw_vocab): flat = {} for key, value in raw_vocab.items(): @@ -32,10 +43,6 @@ def flatten_vocab(raw_vocab): REAL_VOCAB_SIZE = max(vocab_mapping.values()) + 1 # Rule labels/tokens, derived from vocab's rule_tokens, ordered by ID. -# Used only to translate a dataset row's existing rule_ids INDEX (0-12) -# back into its real RULE:xxx vocab token string, so it can be prepended -# to the target sequence. problem_generator.py / rule_ids format is -# unchanged -- this translation happens here in train.py only. _rule_items = sorted(_raw_vocab.get("rule_tokens", {}).items(), key=lambda kv: kv[1]) RULE_TOKEN_STRINGS = [name for name, _ in _rule_items] # e.g. "RULE:power_rule" @@ -81,10 +88,6 @@ def _tokenize(self, envelope, extra_prefix_tokens=None, add_boundaries=False): def __getitem__(self, idx): item = self.data[idx] - # Translate this row's rule_ids index into its real RULE:xxx token - # string, so it becomes part of the sequence the decoder learns to - # generate (as the very first token after [BOS]), instead of a - # separate classifier target. rule_idx = item["rule_ids"] rule_token = ( RULE_TOKEN_STRINGS[rule_idx] @@ -112,15 +115,14 @@ def evaluate_validation(model, val_loader, criterion): model.eval() total_loss = 0.0 total_correct_seq = 0 + total_correct_tokens = 0 + total_valid_tokens = 0 total_seq = 0 steps = 0 + with torch.no_grad(): for batch in val_loader: src_seq = batch["src_seq"] - # Standard teacher-forced shift: tgt_in_seq is tgt_out_seq minus - # the last token; loss is computed against tgt_out_seq minus the - # first token ([BOS]). Both already have [BOS]/[EOS] baked in - # from _tokenize's add_boundaries=True. tgt_in = batch["tgt_in_seq"][:, :-1] tgt_out = batch["tgt_out_seq"][:, 1:] @@ -132,52 +134,71 @@ def evaluate_validation(model, val_loader, criterion): preds = logits.argmax(dim=-1) mask = tgt_out != PAD_ID - correct = ((preds == tgt_out) | ~mask).all(dim=1) - total_correct_seq += correct.sum().item() + + # Per-token accuracy logic + correct_token_mask = (preds == tgt_out) & mask + total_correct_tokens += correct_token_mask.sum().item() + total_valid_tokens += mask.sum().item() + + # Exact sequence match logic + correct_seq = ((preds == tgt_out) | ~mask).all(dim=1) + total_correct_seq += correct_seq.sum().item() total_seq += tgt_out.size(0) steps += 1 if steps == 0: - return 0.0, 0.0 - return total_loss / steps, total_correct_seq / max(total_seq, 1) + return 0.0, 0.0, 0.0 + + avg_loss = total_loss / steps + seq_acc = total_correct_seq / max(total_seq, 1) + token_acc = ( + total_correct_tokens / max(total_valid_tokens, 1) + if total_valid_tokens > 0 + else 0.0 + ) + return avg_loss, seq_acc, token_acc -def write_training_results(metrics_log, best_val_loss): + +def write_training_results(metrics_log, best_val_loss, git_commit_hash): docs_dir = Path("docs") docs_dir.mkdir(exist_ok=True) lines = [ "# Training Results", "", + f"**Git Commit Hash:** `{git_commit_hash}`", f"**Best Validation Loss:** {best_val_loss:.4f}" if best_val_loss < float("inf") else "**Best Validation Loss:** N/A", f"**Total Epochs Run:** {len(metrics_log)}", "", "## Per-Epoch Metrics", "", - "| Epoch | Train Loss | Val Loss | Val Seq Accuracy | Checkpoint Saved |", - "|-------|-----------|----------|-------------------|-----------------|", + "| Epoch | Train Loss | Val Loss | Per-Token Acc | Val Seq Acc | Saved |", + "|-------|-----------|----------|---------------|-------------|-------|", ] for m in metrics_log: val_loss = f"{m['val_loss']:.4f}" if m['val_loss'] is not None else "N/A" + token_acc = f"{m['val_token_acc']:.4f}" if m['val_token_acc'] is not None else "N/A" val_acc = f"{m['val_seq_acc']:.4f}" if m['val_seq_acc'] is not None else "N/A" saved = "Yes" if m['saved'] else "No" lines.append( - f"| {m['epoch']} | {m['train_loss']:.4f} | {val_loss} | {val_acc} | {saved} |" + f"| {m['epoch']} | {m['train_loss']:.4f} | {val_loss} | {token_acc} | {val_acc} | {saved} |" ) lines.extend([ "", - "## Configuration", + "## Configuration Snapshot", "", f"- **Architecture:** SimpleCalculusModel (standard nn.Transformer encoder-decoder)", f"- **Learning Rate:** {config.get('learning_rate')}", + f"- **Warmup Steps:** {config.get('warmup_steps', 1000)}", f"- **Batch Size:** {config.get('batch_size')}", f"- **Hidden Dim:** {config.get('hidden_dim')}", f"- **Max Steps/Epoch:** {config.get('max_steps')}", f"- **Early Stopping:** patience={config.get('early_stopping', {}).get('patience', 'N/A')}, min_delta={config.get('early_stopping', {}).get('min_delta', 'N/A')}", f"- **Vocab Size:** {REAL_VOCAB_SIZE}", f"- **Gradient Clipping:** max_norm={config.get('grad_clip_max_norm', 1.0)}", - f"- **Rule prediction:** folded into output sequence as leading RULE:xxx token (see docs/KNOWN_ISSUES.md)", + f"- **Rule Prediction:** Folded into output sequence as leading RULE:xxx token", "", ]) @@ -187,11 +208,12 @@ def write_training_results(metrics_log, best_val_loss): def run_training_pipeline(): - print(f"--- Training SimpleCalculusModel (vocab size: {REAL_VOCAB_SIZE}) ---") + commit_hash = get_git_commit_hash() + print(f"--- Training SimpleCalculusModel (commit: {commit_hash}, vocab: {REAL_VOCAB_SIZE}) ---") train_file = Path("data/splits/train.jsonl") if not train_file.exists(): - print("Train split missing!") + print("CRITICAL: Train split missing! Run problem_generator.py first.") sys.exit(1) train_loader = DataLoader(SlangDatasetLoader(train_file), batch_size=config["batch_size"], shuffle=True) @@ -207,9 +229,20 @@ def run_training_pipeline(): pad_id=PAD_ID, max_len=MAX_LEN, ) - optimizer = torch.optim.Adam(model.parameters(), lr=config["learning_rate"]) - grad_clip_max_norm = config.get("grad_clip_max_norm", 1.0) + base_lr = config["learning_rate"] + optimizer = torch.optim.Adam(model.parameters(), lr=base_lr) + + # Linear Warmup Scheduler setup + warmup_steps = config.get("warmup_steps", 1000) + def lr_lambda(current_step): + if current_step < warmup_steps: + return float(current_step) / float(max(1, warmup_steps)) + return 1.0 + + scheduler = LambdaLR(optimizer, lr_lambda=lr_lambda) + + grad_clip_max_norm = config.get("grad_clip_max_norm", 1.0) criterion = nn.CrossEntropyLoss(ignore_index=PAD_ID) best_val_loss = float("inf") @@ -233,17 +266,7 @@ def run_training_pipeline(): use_early_stopping = False epochs = config.get("epochs", 1) - - if FINAL_CHECKPOINT_PATH.exists(): - try: - model.load_state_dict(torch.load(str(FINAL_CHECKPOINT_PATH), map_location="cpu")) - print(f"Loaded existing checkpoint from {FINAL_CHECKPOINT_PATH} to resume training.") - if val_loader is not None: - val_loss, val_acc = evaluate_validation(model, val_loader, criterion) - best_val_loss = val_loss - print(f"Initial val loss from resumed checkpoint: {best_val_loss:.4f}") - except Exception as e: - print(f"Could not load checkpoint to resume: {e}") + global_step = 0 for epoch in range(1, epochs + 1): model.train() @@ -262,29 +285,38 @@ def run_training_pipeline(): logits = model(src_seq, tgt_in) loss = criterion(logits.reshape(-1, REAL_VOCAB_SIZE), tgt_out.reshape(-1)) loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=grad_clip_max_norm) optimizer.step() + scheduler.step() + global_step += 1 epoch_loss += loss.item() steps_run += 1 avg_train_loss = epoch_loss / max(steps_run, 1) - print(f"Epoch {epoch}/{epochs} - Train Loss: {avg_train_loss:.4f}") + current_lr = scheduler.get_last_lr()[0] + print(f"Epoch {epoch}/{epochs} - Train Loss: {avg_train_loss:.4f} (LR: {current_lr:.6f})") epoch_metrics = { "epoch": epoch, "train_loss": avg_train_loss, "val_loss": None, + "val_token_acc": None, "val_seq_acc": None, "saved": False, } if val_loader is not None: - val_loss, val_acc = evaluate_validation(model, val_loader, criterion) - print(f"Epoch {epoch} - Val Loss: {val_loss:.4f} Val Seq Accuracy: {val_acc:.4f}") + val_loss, val_seq_acc, val_token_acc = evaluate_validation(model, val_loader, criterion) + print( + f"Epoch {epoch} - Val Loss: {val_loss:.4f} | " + f"Token Acc: {val_token_acc:.4f} | Seq Acc: {val_seq_acc:.4f}" + ) epoch_metrics["val_loss"] = val_loss - epoch_metrics["val_seq_acc"] = val_acc + epoch_metrics["val_token_acc"] = val_token_acc + epoch_metrics["val_seq_acc"] = val_seq_acc if val_loss < best_val_loss - (min_delta if use_early_stopping else 0): best_val_loss = val_loss @@ -295,7 +327,7 @@ def run_training_pipeline(): epoch_metrics["saved"] = True else: patience_counter += 1 - print(f" Epoch {epoch}: val loss {val_loss:.4f} did not improve from {best_val_loss:.4f}, skipping checkpoint save.") + print(f" Epoch {epoch}: val loss {val_loss:.4f} did not improve from {best_val_loss:.4f}.") if use_early_stopping and patience_counter >= patience: print("Early stopping triggered. Training stopped.") metrics_log.append(epoch_metrics) @@ -308,7 +340,7 @@ def run_training_pipeline(): metrics_log.append(epoch_metrics) - write_training_results(metrics_log, best_val_loss) + write_training_results(metrics_log, best_val_loss, commit_hash) print("--- Training complete ---") From 09c6aaac82c528209c78592e900125de8b142e79 Mon Sep 17 00:00:00 2001 From: chaudhryumer Date: Thu, 13 Aug 2026 06:45:15 -0700 Subject: [PATCH 2/2] fix the minor errors in multiple files expeciallt tangent_line error --- config.json | 4 ++-- docs/EVAL_RESULTS.md | 8 ++++---- docs/TRAINING_RESULTS.md | 36 ++++++++++++++++++----------------- docs/runs/RUN_LOG.md | 0 inference/verifier.py | 7 ++++++- model/simple_transformer.py | 9 ++++++--- problem_generator.py | 33 +++++++++++++++++++++----------- tokenizer/slang_serializer.py | 5 ++++- 8 files changed, 63 insertions(+), 39 deletions(-) create mode 100644 docs/runs/RUN_LOG.md diff --git a/config.json b/config.json index 0b04d6e..e9611cd 100644 --- a/config.json +++ b/config.json @@ -5,10 +5,10 @@ "max_steps": 3500, "hidden_dim": 256, "max_len": 32, - "epochs": 10, + "epochs": 20, "grad_clip_max_norm": 1.0, "early_stopping": { - "patience": 12, + "patience": 15, "min_delta": 0.0002 }, "validation_logging": true diff --git a/docs/EVAL_RESULTS.md b/docs/EVAL_RESULTS.md index 77c6b6c..e08bdd2 100644 --- a/docs/EVAL_RESULTS.md +++ b/docs/EVAL_RESULTS.md @@ -4,9 +4,9 @@ | Operation | Total Problems | Exact Match (Accuracy) | Verification Rate | |---|---|---|---| -| diff | 80 | 29/80 (36.2%) | 29/80 (36.2%) | +| diff | 80 | 19/80 (23.8%) | 19/80 (23.8%) | | gradient | 50 | 0/50 (0.0%) | 0/50 (0.0%) | -| integrate | 60 | 40/60 (66.7%) | 40/60 (66.7%) | -| partial | 60 | 17/60 (28.3%) | 17/60 (28.3%) | +| integrate | 60 | 37/60 (61.7%) | 37/60 (61.7%) | +| partial | 60 | 9/60 (15.0%) | 9/60 (15.0%) | | tangent_line | 50 | 0/50 (0.0%) | 0/50 (0.0%) | -| **Overall** | **300** | **86/300 (28.7%)** | **86/300 (28.7%)** | +| **Overall** | **300** | **65/300 (21.7%)** | **65/300 (21.7%)** | diff --git a/docs/TRAINING_RESULTS.md b/docs/TRAINING_RESULTS.md index de72691..8564b85 100644 --- a/docs/TRAINING_RESULTS.md +++ b/docs/TRAINING_RESULTS.md @@ -1,31 +1,33 @@ # Training Results -**Best Validation Loss:** 0.0317 +**Git Commit Hash:** `d26bbe19f4b0b694d6bca560b51cd60e42a36cbc` +**Best Validation Loss:** 0.0167 **Total Epochs Run:** 10 ## Per-Epoch Metrics -| Epoch | Train Loss | Val Loss | Val Seq Accuracy | Checkpoint Saved | -|-------|-----------|----------|-------------------|-----------------| -| 1 | 0.0981 | 0.0322 | 0.7711 | Yes | -| 2 | 0.0330 | 0.0317 | 0.7730 | Yes | -| 3 | 0.0323 | 0.0316 | 0.7730 | No | -| 4 | 0.0324 | 0.0318 | 0.7705 | No | -| 5 | 0.0324 | 0.0317 | 0.7730 | No | -| 6 | 0.0317 | 0.0316 | 0.7730 | No | -| 7 | 0.0323 | 0.0316 | 0.7730 | No | -| 8 | 0.0318 | 0.0332 | 0.7676 | No | -| 9 | 0.0320 | 0.0315 | 0.7730 | No | -| 10 | 0.0318 | 0.0314 | 0.7730 | No | +| Epoch | Train Loss | Val Loss | Per-Token Acc | Val Seq Acc | Saved | +|-------|-----------|----------|---------------|-------------|-------| +| 1 | 0.2888 | 0.0200 | 0.9856 | 0.7683 | Yes | +| 2 | 0.0191 | 0.0170 | 0.9898 | 0.8346 | Yes | +| 3 | 0.0179 | 0.0170 | 0.9898 | 0.8345 | No | +| 4 | 0.0176 | 0.0202 | 0.9892 | 0.8252 | No | +| 5 | 0.0173 | 0.0167 | 0.9899 | 0.8363 | Yes | +| 6 | 0.0172 | 0.0171 | 0.9898 | 0.8346 | No | +| 7 | 0.0170 | 0.0168 | 0.9899 | 0.8363 | No | +| 8 | 0.0170 | 0.0173 | 0.9897 | 0.8335 | No | +| 9 | 0.0171 | 0.0167 | 0.9898 | 0.8356 | No | +| 10 | 0.0168 | 0.0169 | 0.9898 | 0.8356 | No | -## Configuration +## Configuration Snapshot - **Architecture:** SimpleCalculusModel (standard nn.Transformer encoder-decoder) - **Learning Rate:** 0.0001 +- **Warmup Steps:** 1000 - **Batch Size:** 32 - **Hidden Dim:** 256 - **Max Steps/Epoch:** 3500 -- **Early Stopping:** patience=8, min_delta=0.0005 -- **Vocab Size:** 106 +- **Early Stopping:** patience=12, min_delta=0.0002 +- **Vocab Size:** 124 - **Gradient Clipping:** max_norm=1.0 -- **Rule prediction:** folded into output sequence as leading RULE:xxx token (see docs/KNOWN_ISSUES.md) +- **Rule Prediction:** Folded into output sequence as leading RULE:xxx token diff --git a/docs/runs/RUN_LOG.md b/docs/runs/RUN_LOG.md new file mode 100644 index 0000000..e69de29 diff --git a/inference/verifier.py b/inference/verifier.py index 13f727a..aa0710f 100644 --- a/inference/verifier.py +++ b/inference/verifier.py @@ -359,7 +359,12 @@ def def_int_fn(inp): elif op == "gradient": oracle_fn = lambda inp: gradient_oracle(inp["expr"], get_variables(inp)) elif op == "tangent_line": - oracle_fn = lambda inp: tangent_line_oracle(inp["expr"], inp["var"], float(inp["point"])) + def _tangent_line_fn(inp): + point_val = inp["point"] + if isinstance(point_val, dict): + point_val = point_val.get(inp["var"], next(iter(point_val.values()))) + return tangent_line_oracle(inp["expr"], inp["var"], float(point_val)) + oracle_fn = _tangent_line_fn elif op == "product_rule": oracle_fn = lambda inp: product_rule_differentiate(inp["u"], inp["v"], inp["var"]) elif op == "quotient_rule": diff --git a/model/simple_transformer.py b/model/simple_transformer.py index 8c5684f..0174ceb 100644 --- a/model/simple_transformer.py +++ b/model/simple_transformer.py @@ -69,9 +69,12 @@ def __init__( @staticmethod def _causal_mask(seq_len, device): - return torch.triu( - torch.full((seq_len, seq_len), float("-inf"), device=device), diagonal=1 - ) + # Bool mask instead of float -inf mask, matching the dtype of the + # padding masks to eliminate PyTorch's mismatched-mask-type warning. + mask = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1) + return mask + + def forward(self, src_seq, tgt_in_seq): device = src_seq.device diff --git a/problem_generator.py b/problem_generator.py index 1144b7f..52d5b7d 100644 --- a/problem_generator.py +++ b/problem_generator.py @@ -291,12 +291,21 @@ def generate_tangent_line_diff(var="x"): ans_terms.append({"coeff": int(intercept)}) ans = {"numi": {"terms": ans_terms}, "deno": 1} - return src, ans, x0, 0 # rule_id 0 = power_rule + # ADDED HERE (1): Wrap expression and point into the tangent_line operation + src_op = {"op": "tangent_line", "var": var, "expr": src, "point": {var: x0}} + + return src_op, ans, x0, 0 # rule_id 0 = power_rule # Fallback: f(x)=x^2 at x0=1 -> tangent line y = 2x - 1 + fallback_src = {"numi": {"terms": [{"coeff": 1, "var": {var: 2}}]}, "deno": 1} + fallback_ans = {"numi": {"terms": [{"coeff": 2, "var": {var: 1}}, {"coeff": -1}]}, "deno": 1} + + # ADDED HERE (2): Wrap fallback src as well so output structure remains consistent + fallback_src_op = {"op": "tangent_line", "var": var, "expr": fallback_src, "point": {var: 1}} + return ( - {"numi": {"terms": [{"coeff": 1, "var": {var: 2}}]}, "deno": 1}, - {"numi": {"terms": [{"coeff": 2, "var": {var: 1}}, {"coeff": -1}]}, "deno": 1}, + fallback_src_op, + fallback_ans, 1, 0, ) @@ -458,8 +467,10 @@ def generate_slang_dataset(): "verification_state": 1, }) - # 12. Gradient (10k) - for _ in range(10000): + # 12. Gradient (30k, increased from 10k — model was not learning the + # NODE:GRADIENT output structure at 10k rows / ~6% of dataset) + # 12. Gradient (30k rows) + for _ in range(30000): expr, ans, rule_id = generate_gradient_diff() src_op = {"op": "gradient", "var": "x", "expr": expr} dataset.append({ @@ -468,21 +479,21 @@ def generate_slang_dataset(): "tgt_output_tokens": ans, "rule_ids": rule_id, "verification_state": 1, - }) + }) + # 13. Tangent line (10k) # 13. Tangent line (10k) for _ in range(10000): var = random.choice(VARIABLES[:1]) - src, ans, x0, rule_id = generate_tangent_line_diff(var) - src_op = {"op": "tangent_line", "var": var, "expr": src, "point": x0} + # generate_tangent_line_diff returns (src_op, ans, x0, rule_id) + src_op, ans, _, rule_id = generate_tangent_line_diff(var) dataset.append({ - "src_tokens": src_op, + "src_tokens": src_op, # Use src_op directly! "tgt_input_tokens": ans, "tgt_output_tokens": ans, "rule_ids": rule_id, "verification_state": 1, - }) - + }) random.shuffle(dataset) with open("data/slang_dataset.jsonl", "w", encoding="utf-8") as f: diff --git a/tokenizer/slang_serializer.py b/tokenizer/slang_serializer.py index 5d9eb33..43771a8 100644 --- a/tokenizer/slang_serializer.py +++ b/tokenizer/slang_serializer.py @@ -102,7 +102,10 @@ def serialize_op_node(n: Dict[str, Any]) -> None: # to keep parse_op_node's fixed decorator order unambiguous. Only # tangent_line sets this field; all other op-nodes are unaffected. if "point" in n: - point_val = float(n["point"]) + point_raw = n["point"] + if isinstance(point_raw, dict): + point_raw = point_raw.get(n.get("var"), next(iter(point_raw.values()))) + point_val = float(point_raw) if point_val.is_integer(): point_val = int(point_val) tokens.append(f"{POINT_PREFIX}{point_val}")