| """Frozen R51 generation path with unused adapter branches removed. |
| |
| No model hooks or trained modules; see source_manifest.json and parity tests. |
| """ |
|
|
| import math |
| import time |
|
|
| import torch |
| from transformers import TemperatureLogitsWarper, TopKLogitsWarper, TopPLogitsWarper |
| from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS |
|
|
|
|
| def new_cache(model): |
| from transformers.models.granitemoehybrid.modeling_granitemoehybrid import ( |
| HybridMambaAttentionDynamicCache, |
| ) |
|
|
| parameter = next(model.parameters()) |
| cache = HybridMambaAttentionDynamicCache( |
| model.config, batch_size=1, dtype=parameter.dtype, device=parameter.device |
| ) |
| for layer in cache.transformer_layers: |
| shape = ( |
| 1, |
| model.config.num_key_value_heads, |
| 0, |
| model.config.hidden_size // model.config.num_attention_heads, |
| ) |
| cache.key_cache[layer] = torch.empty(shape, dtype=parameter.dtype, device=parameter.device) |
| cache.value_cache[layer] = torch.empty_like(cache.key_cache[layer]) |
| return cache |
|
|
|
|
| def idle_frozen(model): |
| if model.training or any(p.requires_grad or p.grad is not None for p in model.parameters()): |
| raise ValueError("Runtime requires a frozen evaluation backbone") |
| attention = ALL_ATTENTION_FUNCTIONS["sdpa"] |
| if ( |
| getattr(model, "_jev_feedback_active", False) |
| or getattr(attention, "_r17_controlled", False) |
| or getattr(attention, "_r18_controlled", False) |
| or any(m._forward_hooks or m._forward_pre_hooks for m in model.modules()) |
| or torch.nn.modules.module._global_forward_hooks |
| or torch.nn.modules.module._global_forward_pre_hooks |
| ): |
| raise RuntimeError("Model intervention already active") |
|
|
|
|
| def valid_ids(model, ids): |
| return ( |
| isinstance(ids, list) |
| and bool(ids) |
| and all(type(i) is int and 0 <= i < model.config.vocab_size for i in ids) |
| ) |
|
|
|
|
| def sample_setup(sampling, device): |
| if sampling is None: |
| return (None, [], None) |
| if ( |
| not isinstance(sampling, dict) |
| or set(sampling) != {"seed", "temperature", "top_p"} |
| or type(sampling["seed"]) is not int |
| or (not 0 <= sampling["seed"] < 2**63) |
| or any( |
| type(sampling[k]) not in (int, float) or not math.isfinite(sampling[k]) |
| for k in ("temperature", "top_p") |
| ) |
| or (not 0 < sampling["temperature"] <= 2) |
| or (not 0 < sampling["top_p"] <= 1) |
| ): |
| raise ValueError("Invalid sampling settings") |
| settings = dict(sampling, top_k=50) |
| generator = torch.Generator(device=device).manual_seed(settings["seed"]) |
| warpers = [ |
| TemperatureLogitsWarper(float(settings["temperature"])), |
| TopKLogitsWarper(50), |
| TopPLogitsWarper(float(settings["top_p"])), |
| ] |
| return (generator, warpers, settings) |
|
|
|
|
| def generate(model, tok, ids, *, limit, eos, context_limit=32768, sampling=None, deadline=None): |
| idle_frozen(model) |
| if ( |
| not valid_ids(model, ids) |
| or not valid_ids(model, eos) |
| or len(eos) != len(set(eos)) |
| or (type(limit) is not int) |
| or (not 1 <= limit <= 2048) |
| or (type(context_limit) is not int) |
| or (not 1 <= context_limit <= 32768) |
| or (context_limit > model.config.max_position_embeddings) |
| or (len(ids) + limit > context_limit) |
| ): |
| raise ValueError("Invalid tokens, generation limit or context") |
| device = next(model.parameters()).device |
| generator, warpers, settings = sample_setup(sampling, device) |
| pending, accepted, past = (list(ids), [], 0) |
| cache, traces = (new_cache(model), []) |
| tick = time.monotonic() |
| forward_pending, attempted_positions = False, 0 |
| decoding_started = False |
| try: |
| with torch.no_grad(): |
| for _ in range(limit): |
| if deadline: |
| deadline() |
| attempted_positions += len(pending) |
| forward_pending = True |
| out = model( |
| input_ids=torch.tensor([pending], device=device), |
| attention_mask=torch.ones( |
| (1, past + len(pending)), device=device, dtype=torch.long |
| ), |
| past_key_values=cache, |
| use_cache=True, |
| cache_position=torch.arange(past, past + len(pending), device=device), |
| logits_to_keep=1, |
| ) |
| past += len(pending) |
| forward_pending = False |
| cache = out.past_key_values |
| logits = out.logits[:, -1, :].float() |
| if not torch.isfinite(logits).all(): |
| raise ValueError("Nonfinite generation logits") |
| if generator is None: |
| token = int(logits[0].argmax()) |
| else: |
| for warper in warpers: |
| logits = warper(None, logits) |
| token = int( |
| torch.multinomial(torch.softmax(logits, dim=-1), 1, generator=generator)[ |
| 0, 0 |
| ] |
| ) |
| accepted.append(token) |
| if token in eos: |
| break |
| pending = [token] |
| decoding_started = True |
| body = accepted[:-1] if accepted[-1] in eos else accepted |
| return dict( |
| prompt_token_ids=list(ids), |
| generated_token_ids=accepted, |
| text=tok.decode(body, skip_special_tokens=False, clean_up_tokenization_spaces=False), |
| finish_reason="eos" if accepted[-1] in eos else "length", |
| processed_tokens=len(ids) + len(accepted) - 1, |
| seconds=time.monotonic() - tick, |
| events=traces, |
| sampling=settings, |
| ) |
| except BaseException as exc: |
| body = accepted[:-1] if accepted and accepted[-1] in eos else accepted |
| partial = dict( |
| prompt_token_ids=list(ids), |
| generated_token_ids=list(accepted), |
| text=None, |
| finish_reason=( |
| "time" |
| if isinstance(exc, TimeoutError) |
| else "error" |
| if isinstance(exc, Exception) |
| else "interrupted" |
| ), |
| processed_tokens=past, |
| attempted_forward_positions=attempted_positions, |
| forward_work_unknown=forward_pending, |
| seconds=time.monotonic() - tick, |
| events=traces, |
| sampling=settings, |
| ) |
| exc.partial_generation = partial |
| if not decoding_started: |
| try: |
| partial["text"] = tok.decode( |
| body, skip_special_tokens=False, clean_up_tokenization_spaces=False |
| ) |
| except BaseException as decoding_error: |
| partial["decoding_error_type"] = type(decoding_error).__name__ |
| if not isinstance(decoding_error, Exception): |
| decoding_error.partial_generation = partial |
| raise |
| raise |
|
|
