Spaces:
Running on Zero
Running on Zero
| """Opt-in selector measurements; never runs an additional model forward pass.""" | |
| import time | |
| import torch | |
| class ProfiledBigramSelector: | |
| """Wrap a selector to measure GPU-completed time and same-logit argmax changes. | |
| Explicit synchronization and metric collection add benchmark overhead. This | |
| wrapper is not used during normal generation. Frozen positions are excluded | |
| from the change fraction; each editable prediction at each step is counted. | |
| """ | |
| def __init__(self, selector): | |
| self.selector = selector | |
| self.calls = self.changed = self.editable = 0 | |
| self.elapsed_ms = 0.0 | |
| def set_step_context(self, block_step, block_steps): | |
| setter = getattr(self.selector, 'set_step_context', None) | |
| if setter is not None: | |
| setter(block_step, block_steps) | |
| def select_block(self, logits, current_ids, editable, **kwargs): | |
| if logits.device.type == 'cuda': | |
| torch.cuda.synchronize(logits.device) | |
| start = time.perf_counter() | |
| result = self.selector.select_block(logits, current_ids, editable, **kwargs) | |
| if logits.device.type == 'cuda': | |
| torch.cuda.synchronize(logits.device) | |
| self.elapsed_ms += (time.perf_counter() - start) * 1000 | |
| self.calls += 1 | |
| self.changed += int(((result[0] != logits.argmax(-1)) & editable).sum()) | |
| self.editable += int(editable.sum()) | |
| return result | |
| def last_path_margins(self): | |
| return self.selector.last_path_margins | |
| def metrics(self): | |
| return dict(selector_calls=self.calls, editable_predictions=self.editable, | |
| changed_predictions=self.changed, | |
| changed_fraction=self.changed / max(1, self.editable), | |
| selector_ms=self.elapsed_ms, | |
| mean_selector_ms=self.elapsed_ms / max(1, self.calls)) | |