invarlock 0.3.2__py3-none-any.whl → 0.3.4__py3-none-any.whl

invarlock/eval/data.py

@@ -950,6 +950,9 @@ class WikiText2Provider:
             if override_size is not None:
                 batch_size = max(1, min(override_size, len(candidates)))
 
+            config = getattr(model, "config", None)
+            scorer_vocab_size = getattr(config, "vocab_size", None)
+
             input_batch: list[list[int]] = []
             attention_batch: list[list[int]] = []
             candidate_batch: list[dict[str, Any]] = []
@@ -970,6 +973,14 @@ class WikiText2Provider:
                             attention_batch, dtype=torch.long, device=device
                         )
 
+                        # Guard against out-of-range token IDs when scoring with GPT-2.
+                        # Some model tokenizers emit IDs beyond GPT-2 vocab, which can
+                        # trigger device-side asserts in embedding/gather kernels.
+                        if scorer_vocab_size and scorer_vocab_size > 0:
+                            input_tensor = input_tensor.clamp(
+                                min=0, max=scorer_vocab_size - 1
+                            )
+
                         outputs = model(input_tensor, attention_mask=attention_tensor)
                         shift_logits = outputs.logits[:, :-1, :].contiguous()
                         shift_labels = input_tensor[:, 1:].contiguous()

invarlock/eval/primary_metric.py

@@ -214,9 +214,15 @@ class _PPLCausal(PrimaryMetric):
     ) -> dict[str, Any]:
         subj = self._coerce_contrib_array(subject)
         base = self._coerce_contrib_array(baseline)
-        # Compute simple (unweighted) per-example arrays in log space; weights ignored for bootstrap here
+        # Compute per-example arrays in log space; use weights for paired bootstrap
         subj_vals = [v for (v, _w) in subj]
         base_vals = [v for (v, _w) in base]
+        pair_weights = []
+        for (_sv, sw), (_bv, bw) in zip(subj, base, strict=False):
+            weight = bw if math.isfinite(bw) and bw > 0 else sw
+            if not math.isfinite(weight) or weight <= 0:
+                weight = 1.0
+            pair_weights.append(float(weight))
 
         # Points in display space
         def _point(
@@ -249,15 +255,24 @@ class _PPLCausal(PrimaryMetric):
         dlog_lo, dlog_hi = compute_paired_delta_log_ci(
             subj_vals,
             base_vals,
+            weights=pair_weights,
             method="bca",
             replicates=reps_eff,
             alpha=alpha,
             seed=seed_eff,
         )
-        delta_log = float(
-            sum((s - b) for s, b in zip(subj_vals, base_vals, strict=False))
-            / max(1, min(len(subj_vals), len(base_vals)))
-        )
+        if pair_weights and len(pair_weights) >= min(len(subj_vals), len(base_vals)):
+            sw = 0.0
+            swx = 0.0
+            for s, b, w in zip(subj_vals, base_vals, pair_weights, strict=False):
+                sw += w
+                swx += w * (s - b)
+            delta_log = float(swx / sw) if sw > 0 else float("nan")
+        else:
+            delta_log = float(
+                sum((s - b) for s, b in zip(subj_vals, base_vals, strict=False))
+                / max(1, min(len(subj_vals), len(base_vals)))
+            )
         ratio = self.display_transform(delta_log)
         return {
             "kind": self.kind,