early_exit_wrapper.py

# early_exit_wrapper.py
import json, torch, torch.nn as nn, torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

class EarlyExitClassifier(nn.Module):
    def __init__(self, hidden_size, vocab_size, dtype=torch.float16, device=None):
        super().__init__()
        self.linear = nn.Linear(hidden_size, vocab_size, dtype=dtype)
        if device is not None:
            self.to(device)

    def forward(self, hidden_states):
        return self.linear(hidden_states)

class EarlyExitModelWrapper(nn.Module):
    def __init__(self, model, confidence_threshold=0.9, num_layers_to_check=None, device=None, dtype=torch.float16):
        super().__init__()
        self.model = model
        self.config = model.config
        self.confidence_threshold = confidence_threshold
        nl = num_layers_to_check or len(model.model.layers)
        self.classifiers = nn.ModuleList([
            EarlyExitClassifier(self.config.hidden_size, self.config.vocab_size, dtype=dtype, device=device)
            for _ in range(nl)
        ])
        self._device = device or next(model.parameters()).device
        self._dtype = dtype

    def forward(self, input_ids, attention_mask=None, **kwargs):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True, **kwargs)
        hidden_states = outputs.hidden_states
        for layer_idx, classifier in enumerate(self.classifiers, start=1):
            h_state = hidden_states[layer_idx]
            logits = classifier(h_state[:, -1, :].to(self._dtype))
            probs = F.softmax(logits, dim=-1)
            max_prob = probs.max(dim=-1)[0]
            if bool((max_prob >= self.confidence_threshold).item()):
                return {"logits": logits, "exit_layer": layer_idx, "hidden_states": h_state}

        final_hidden = self.model.model.norm(hidden_states[-1].to(self._dtype))
        final_logits = self.model.lm_head(final_hidden[:, -1, :])
        return {"logits": final_logits, "exit_layer": len(hidden_states)-1, "hidden_states": hidden_states[-1]}

@torch.no_grad()
def generate_with_early_exit(prompt, model, tokenizer, max_new_tokens=64, temperature=0.7, top_p=0.9, device=None):
    device = device or next(model.parameters()).device
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    input_ids = inputs["input_ids"]
    attention_mask = inputs.get("attention_mask", None)
    generated_ids = input_ids.clone()

    for _ in range(max_new_tokens):
        outputs = model(input_ids=generated_ids, attention_mask=attention_mask)
        logits = outputs["logits"] / temperature

        if top_p < 1.0:
            sorted_logits, sorted_indices = torch.sort(logits, dim=-1, descending=True)
            cumprobs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
            to_remove = cumprobs > top_p
            to_remove[:, 1:] = to_remove[:, :-1].clone()
            to_remove[:, 0] = 0
            indices_to_remove = to_remove.scatter(1, sorted_indices, to_remove)
            logits[indices_to_remove] = float("-inf")

        probs = torch.softmax(logits, dim=-1)
        next_token_id = torch.multinomial(probs, num_samples=1)
        generated_ids = torch.cat([generated_ids, next_token_id], dim=-1)
        if attention_mask is not None:
            attention_mask = torch.cat([attention_mask, torch.ones_like(next_token_id)], dim=-1)
        if next_token_id.item() == tokenizer.eos_token_id:
            break
    return generated_ids

def load_early_exit_from_hub(repo_id: str, device: str = None):
    """
    Loads:
      - early_exit_config.json
      - early_exit_heads.safetensors
    and returns (wrapped_model, tokenizer).
    """
    cfg_path = hf_hub_download(repo_id=repo_id, filename="early_exit_config.json")
    with open(cfg_path, "r") as f:
        cfg = json.load(f)

    base_id = cfg["base_model"]
    dtype = torch.float16 if cfg.get("dtype", "float16") == "float16" else torch.float32
    device = device or ("mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu"))

    tokenizer = AutoTokenizer.from_pretrained(base_id)
    base_model = AutoModelForCausalLM.from_pretrained(
        base_id,
        torch_dtype=dtype,
        device_map={"": device} if device != "cpu" else None,
    )

    wrapped = EarlyExitModelWrapper(
        base_model,
        confidence_threshold=float(cfg["confidence_threshold"]),
        num_layers_to_check=int(cfg["num_layers_to_check"]),
        device=device,
        dtype=dtype,
    )

    heads_path = hf_hub_download(repo_id=repo_id, filename="early_exit_heads.safetensors")
    state = load_file(heads_path)
    wrapped.classifiers.load_state_dict(state, strict=True)

    return wrapped, tokenizer