Upload 4 files

config.json

@@ -17,6 +17,16 @@
     "longest_word_chars",
     "vocab_size"
   ],
+  "base_feature_names": [
+    "chars",
+    "words",
+    "avg_chars_per_word",
+    "punctuation_ratio",
+    "symbol_ratio",
+    "longest_word_chars",
+    "vocab_size"
+  ],
+  "base_feature_dim": 7,
   "feature_dim": 19,
   "num_layers": 8,
   "hidden_dim": 32,
@@ -27,5 +37,47 @@
   "use_log1p_features": true,
   "use_log1p_target": false,
   "center_target": false,
-  "target_offset": 0.0
+  "target_offset": 0.0,
+  "feature_mean": [
+    4971.99609375,
+    750.8612670898438,
+    5.387211799621582,
+    0.038271043449640274,
+    0.00981982797384262,
+    24.875553131103516,
+    62374.8046875,
+    6.58632230758667,
+    0.15978458523750305,
+    7.920589923858643,
+    6.06857967376709,
+    9.939688682556152,
+    216.70751953125,
+    61.361202239990234,
+    4028.412353515625,
+    29.0120792388916,
+    1.2395837306976318,
+    32.29179382324219,
+    278.0736999511719
+  ],
+  "feature_std": [
+    10841.10546875,
+    1556.9542236328125,
+    1.7252269983291626,
+    0.028954673558473587,
+    0.020614376291632652,
+    540.275390625,
+    76498.015625,
+    2.147888660430908,
+    0.028698621317744255,
+    0.9891781806945801,
+    0.963409960269928,
+    1.7922862768173218,
+    818.908935546875,
+    420.1692199707031,
+    8589.4033203125,
+    90.29910278320312,
+    5.135631084442139,
+    561.3566284179688,
+    1134.063232421875
+  ]
 }

configuration_jetoncount.py

@@ -9,6 +9,8 @@ class JetonCountConfig(PretrainedConfig):
     def __init__(
         self,
         feature_names=None,
+        base_feature_names=None,
+        base_feature_dim=7,
         feature_dim=19,
         num_layers=8,
         hidden_dim=32,
@@ -20,9 +22,21 @@ class JetonCountConfig(PretrainedConfig):
         use_log1p_target=False,
         center_target=False,
         target_offset=0.0,
+        feature_mean=None,
+        feature_std=None,
         **kwargs,
     ):
-        self.feature_names = feature_names or []
+        self.base_feature_names = base_feature_names or [
+            "chars",
+            "words",
+            "avg_chars_per_word",
+            "punctuation_ratio",
+            "symbol_ratio",
+            "longest_word_chars",
+            "vocab_size",
+        ]
+        self.feature_names = feature_names or list(self.base_feature_names)
+        self.base_feature_dim = int(base_feature_dim)
         self.feature_dim = int(feature_dim)
         self.num_layers = int(num_layers)
         self.hidden_dim = int(hidden_dim)
@@ -34,4 +48,6 @@ class JetonCountConfig(PretrainedConfig):
         self.use_log1p_target = bool(use_log1p_target)
         self.center_target = bool(center_target)
         self.target_offset = float(target_offset)
+        self.feature_mean = feature_mean
+        self.feature_std = feature_std
         super().__init__(**kwargs)

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4c048d0728fd46f7bff885866ffdcbe742789ff633eb7fb6468cf11f5fda3a34
-size 29148
+oid sha256:f8871d7515ae031717c953aa4008f2ebd6d5b4d2f04ed8c15b6f61f7db11628c
+size 29276

modeling_jetoncount.py

@@ -1,3 +1,5 @@
+"""JetonCount MLP regression model."""
+
 from __future__ import annotations
 
 from dataclasses import dataclass
@@ -32,8 +34,8 @@ def _get_activation(name: str) -> nn.Module:
 
 def _engineer_features_tensor(base: torch.Tensor) -> torch.Tensor:
     """
-    base must be shape [B, 7] or [7].
-    Output is shape [B, 19] or [19].
+    base shape: [7] or [B, 7]
+    output shape: [19] or [B, 19]
     Matches train_mlp_token_regressor.py.
     """
     squeeze = False
@@ -52,34 +54,20 @@ def _engineer_features_tensor(base: torch.Tensor) -> torch.Tensor:
     vocab_size = base[:, 6]
 
     eps = 1e-6
-    ratio_chars_words = chars / torch.clamp(words, min=1.0)
-    ratio_words_chars = words / torch.clamp(chars, min=1.0)
-    log_chars = torch.log1p(torch.clamp(chars, min=0.0))
-    log_words = torch.log1p(torch.clamp(words, min=0.0))
-    log_vocab = torch.log1p(torch.clamp(vocab_size, min=0.0))
-
-    chars_x_punct = chars * punctuation_ratio
-    chars_x_symbol = chars * symbol_ratio
-    words_x_avg = words * avg_chars_per_word
-    words_x_punct = words * punctuation_ratio
-    longest_x_punct = longest_word_chars * punctuation_ratio
-    complexity_proxy = (avg_chars_per_word + longest_word_chars) * (1.0 + punctuation_ratio + symbol_ratio)
-    density_proxy = (chars + eps) * (punctuation_ratio + symbol_ratio + eps)
-
     extra = torch.stack(
         [
-            ratio_chars_words,
-            ratio_words_chars,
-            log_chars,
-            log_words,
-            log_vocab,
-            chars_x_punct,
-            chars_x_symbol,
-            words_x_avg,
-            words_x_punct,
-            longest_x_punct,
-            complexity_proxy,
-            density_proxy,
+            chars / torch.clamp(words, min=1.0),
+            words / torch.clamp(chars, min=1.0),
+            torch.log1p(torch.clamp(chars, min=0.0)),
+            torch.log1p(torch.clamp(words, min=0.0)),
+            torch.log1p(torch.clamp(vocab_size, min=0.0)),
+            chars * punctuation_ratio,
+            chars * symbol_ratio,
+            words * avg_chars_per_word,
+            words * punctuation_ratio,
+            longest_word_chars * punctuation_ratio,
+            (avg_chars_per_word + longest_word_chars) * (1.0 + punctuation_ratio + symbol_ratio),
+            (chars + eps) * (punctuation_ratio + symbol_ratio + eps),
         ],
         dim=-1,
     )
@@ -130,9 +118,8 @@ class JetonCountForRegression(PreTrainedModel):
         if not self.config.standardize_features:
             return x
 
-        mean = self.config.feature_mean
-        std = self.config.feature_std
-
+        mean = getattr(self.config, "feature_mean", None)
+        std = getattr(self.config, "feature_std", None)
         if mean is None or std is None:
             return x
 
@@ -165,24 +152,21 @@ class JetonCountForRegression(PreTrainedModel):
     def _remap_state_dict_keys(self, state_dict):
         """
         Accepts several historical layouts:
-        - 0.weight / 0.bias
-        - net.0.weight / net.0.bias
-        - mlp.net.0.weight / mlp.net.0.bias
+        - mlp.net.0.weight
+        - net.0.weight
+        - 0.weight
         """
         if not state_dict:
             return state_dict
 
-        if any(k.startswith("mlp.net.") or k.startswith("net.") for k in state_dict):
-            return state_dict
-
         remapped = {}
         for k, v in state_dict.items():
             if k.startswith("mlp.net."):
-                remapped[k[len("mlp.") :]] = v  # -> net.0.weight
-            elif k.startswith("mlp."):
-                remapped[k[len("mlp.") :]] = v
-            elif k[0].isdigit():
-                remapped[f"net.{k}"] = v
+                remapped[k] = v
+            elif k.startswith("net."):
+                remapped[f"mlp.{k}"] = v
+            elif k and k[0].isdigit():
+                remapped[f"mlp.net.{k}"] = v
             else:
                 remapped[k] = v
         return remapped
@@ -210,4 +194,4 @@ class JetonCountForRegression(PreTrainedModel):
             labels = labels.to(logits.dtype)
             loss = torch.nn.functional.mse_loss(logits, labels)
 
-        return RegressionOutput(loss=loss, logits=logits)
+        return RegressionOutput(loss=loss, logits=logits)