update

README.md

@@ -0,0 +1,92 @@
+---
+library_name: transformers
+tags:
+  - custom_generate
+  - sampling
+  - kvcache
+---
+
+# Sampling with KV Cache
+
+## Description
+A clean, hackable implementation of sampling (also called ancestral sampling or multinomial sampling). This is a simplified alternative to the complex generation mixin in transformers, designed for readability and ease of modification while maintaining full performance.
+
+The implementation supports both sampling and greedy decoding modes, with optional temperature scaling and top-k/top-p filtering.
+
+## Base model
+- [HuggingFaceTB/SmolLM2-135M-Instruct](https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct)
+
+## Model compatibility
+Most transformer LLM/VLM models trained for causal language modeling.
+
+## Relevant Arguments
+- `temperature` (float): Sampling temperature (default: 1.0, higher = more random)
+- `top_k` (int): Only consider top-k most probable tokens (default: None)
+- `top_p` (float): Only consider tokens with cumulative probability <= top_p (default: None)
+- `do_sample` (bool): Whether to use sampling (True, default) or greedy decoding (False)
+
+### Logits Processing Order
+Logits processors are applied in sequence: `temperature → softmax → top_k → top_p` (same as HuggingFace's `LogitProcessor` system). Temperature scaling occurs before top-p filtering, affecting the probability distribution that top-p operates on.
+
+For example, with `temperature=1.0`, `top_p=0.9` might include tokens A, B, C. With `temperature=0.5`, probability mass is much more concentrated, so `top_p=0.9` might only include token A.
+
+## Outputs
+When `return_dict_in_generate=True`, returns a dictionary with:
+- `sequences`: Generated token IDs
+- `scores`: Log probabilities of sampled tokens (with temperature/sampling modifications)
+- `logprobs`: Original model log probabilities (T=1, no modifications)
+Otherwise, returns a tensor of generated token IDs.
+
+## Example usage
+
+```py
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", device_map="auto")
+
+inputs = tokenizer(["The quick brown"], return_tensors="pt").to(model.device)
+
+# Basic sampling
+gen_out = model.generate(**inputs, custom_generate="manueldeprada/sampling", trust_remote_code=True)
+
+# With temperature
+gen_out = model.generate(**inputs, custom_generate="manueldeprada/sampling", temperature=0.8, trust_remote_code=True)
+
+# With top-k
+gen_out = model.generate(**inputs, custom_generate="manueldeprada/sampling", top_k=50, trust_remote_code=True)
+
+# With top-p (nucleus sampling)
+gen_out = model.generate(**inputs, custom_generate="manueldeprada/sampling", top_p=0.9, trust_remote_code=True)
+
+# Greedy decoding (no sampling)
+gen_out = model.generate(**inputs, custom_generate="manueldeprada/sampling", do_sample=False, trust_remote_code=True)
+
+# Get detailed output with probabilities
+gen_out = model.generate(
+    **inputs, 
+    custom_generate="manueldeprada/sampling", 
+    return_dict_in_generate=True,
+    trust_remote_code=True
+)
+print(f"Generated text: {tokenizer.batch_decode(gen_out['sequences'], skip_special_tokens=True)}")
+print(f"Sampling scores: {gen_out['scores']}")
+print(f"Model log probabilities: {gen_out['logprobs']}")
+```
+
+## Algorithm
+1. Prepare input sequences
+2. For each generation step:
+   - Get logits from the model for the current sequence
+   - Apply temperature scaling to logits
+   - Optionally apply top-k filtering (keep only top-k tokens)
+   - Optionally apply top-p filtering (nucleus sampling)
+   - Convert to probabilities using softmax
+   - Sample from the probability distribution (or take argmax for greedy)
+   - Append the selected token to the sequence
+   - Track sequence completion
+3. Return generated sequences and probability information
+
+
+
+

config.json

@@ -0,0 +1,29 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "hidden_act": "silu",
+  "hidden_size": 576,
+  "initializer_range": 0.041666666666666664,
+  "intermediate_size": 1536,
+  "is_llama_config": true,
+  "max_position_embeddings": 8192,
+  "model_type": "llama",
+  "num_attention_heads": 9,
+  "num_hidden_layers": 30,
+  "num_key_value_heads": 3,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_interleaved": false,
+  "rope_scaling": null,
+  "rope_theta": 100000,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.40.1",
+  "use_cache": true,
+  "vocab_size": 49152
+}

custom_generate/generate.py

@@ -0,0 +1,272 @@
+import torch
+from transformers import Cache, DynamicCache
+from transformers.generation.utils import ModelOutput
+from typing import Optional, Any
+
+def prepare_inputs_for_generation(
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.LongTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ):
+        input_ids = input_ids[:, cache_position].clone(memory_format=torch.contiguous_format)
+        model_inputs = {"cache_position": cache_position, 
+                        "past_key_values": None, 
+                        "input_ids": input_ids, 
+                        "inputs_embeds": None,
+                        "attention_mask": attention_mask,
+                        }
+        if attention_mask is not None and kwargs.get("position_ids") is None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            kwargs["position_ids"] = position_ids
+        model_inputs.update({k: v for k, v in kwargs.items() if k not in model_inputs})
+        return model_inputs
+
+def update_model_kwargs_for_generation(
+        outputs: ModelOutput,
+        model_kwargs: dict[str, Any],
+        num_new_tokens: int = 1,
+    ) -> dict[str, Any]:
+        if "token_type_ids" in model_kwargs:
+            token_type_ids = model_kwargs["token_type_ids"]
+            model_kwargs["token_type_ids"] = torch.cat([token_type_ids, token_type_ids[:, -1].unsqueeze(-1)], dim=-1)
+        if "attention_mask" in model_kwargs:
+            attention_mask = model_kwargs["attention_mask"]
+            model_kwargs["attention_mask"] = torch.cat(
+                [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
+            )
+        past_positions = model_kwargs.pop("cache_position")
+        new_positions = torch.arange(
+            past_positions[-1] + 1, past_positions[-1] + num_new_tokens + 1, dtype=past_positions.dtype
+        ).to(past_positions.device)
+        model_kwargs["cache_position"] = torch.cat((past_positions, new_positions))
+        return model_kwargs
+
+
+def next_logits_with_cache_update(model, model_kwargs, input_ids):
+    """
+    Gets the next token logits and updates the KV cache:
+    - Runs the model forward pass
+    - Extracts logits for the last token
+    - Updates the KV cache
+    - Returns updated `model_kwargs` and `logits`
+
+    Args:
+        model: The language model
+        model_kwargs: Model keyword arguments including KV cache
+        input_ids: Current input token IDs
+
+    Returns:
+        Updated model_kwargs, logits for the next token
+    """
+    model_inputs = prepare_inputs_for_generation(input_ids, **model_kwargs)
+    with torch.no_grad():
+        outputs = model(**model_inputs, return_dict=True)
+
+    logits = outputs.logits[:, -1].detach()
+    model_kwargs = update_model_kwargs_for_generation(outputs, model_kwargs)
+    del outputs
+    return model_kwargs, logits
+
+
+def init_gen(model_kwargs, model, max_new_tokens, bos_token_id):
+    """
+    Initializes the generation process and prepares the KV cache:
+    - Sets up input sequences and model inputs
+    - Prepares the KV cache for generation
+    - Returns updated `model_kwargs` and `input_ids`
+
+    Args:
+        model_kwargs: Model keyword arguments
+        model: The language model
+        max_new_tokens: Maximum number of new tokens to generate
+        bos_token_id: Beginning-of-sequence token ID
+
+    Returns:
+        Model keyword arguments and input token IDs
+    """
+    input_ids = model_kwargs.pop("input_ids")
+    model_kwargs["past_key_values"] = None
+    model_kwargs["cache_position"] = torch.ones(input_ids.shape[1], dtype=torch.int64, device=input_ids.device).cumsum(0) - 1
+    return model_kwargs, input_ids
+
+
+def _apply_top_k(ps, model):
+    """Apply top-k filtering to probabilities."""
+    if not hasattr(model, "generation_config") or not hasattr(
+        model.generation_config, "top_k"
+    ):
+        return ps
+
+    top_k = model.generation_config.top_k
+    if top_k is None or top_k >= ps.size(-1):
+        return ps
+
+    indices_to_remove = ps < torch.topk(ps, top_k)[0][..., -1, None]
+    ps[indices_to_remove] = 0.0
+    return ps / ps.sum(dim=-1, keepdim=True)
+
+
+def _apply_top_p(ps, model):
+    """Apply top-p (nucleus) filtering to probabilities."""
+    if not hasattr(model, "generation_config") or not hasattr(
+        model.generation_config, "top_p"
+    ):
+        return ps
+
+    top_p = model.generation_config.top_p
+    if top_p is None or top_p >= 1.0:
+        return ps
+
+    sorted_probs, sorted_indices = torch.sort(ps, descending=True)
+    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+
+    sorted_indices_to_remove = cumulative_probs > top_p
+    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+    sorted_indices_to_remove[..., 0] = 0
+
+    indices_to_remove = sorted_indices_to_remove.scatter(
+        1, sorted_indices, sorted_indices_to_remove
+    )
+    ps[indices_to_remove] = 0.0
+    return ps / ps.sum(dim=-1, keepdim=True)
+
+
+def sampling(
+    model_kwargs,
+    model,
+    eos_token_ids,
+    pad_token_id,
+    bos_token_id,
+    do_sample=True,
+    max_new_tokens=20,
+    temperature=1.0,
+):
+    """
+    Sampling implementation with proper KV caching.
+
+    Args:
+        prompts: List of input prompts
+        model: The language model
+        max_new_tokens: Maximum number of new tokens to generate
+        eos_token_ids: List of end-of-sequence token IDs
+        pad_token_id: Padding token ID
+        bos_token_id: Beginning-of-sequence token ID
+        max_new_tokens: Maximum number of new tokens to generate
+
+    Returns:
+        Generated sequences, log probabilities, and metadata
+    """
+    # Initialize the generation process and prepare the KV cache
+    model_kwargs, input_ids = init_gen(
+        model_kwargs, model, max_new_tokens, bos_token_id
+    )
+    batch_size, _ = input_ids.shape
+
+    # Keeps track of which sequences are finished and their lengths
+    active_seqs = input_ids.new_ones((batch_size, 1), dtype=torch.bool)
+    # Modified log probabilities of the sequences
+    scores = torch.zeros((batch_size, max_new_tokens), dtype=model.dtype)
+    # Unfiltered sequence log probabilities (temperature=1, no sampling processors applied)
+    logprobs = torch.zeros((batch_size, max_new_tokens), dtype=model.dtype)
+
+    for i in range(max_new_tokens):
+        # Get the next token probabilities and update the KV cache
+        model_kwargs, logits = next_logits_with_cache_update(
+            model, model_kwargs, input_ids
+        )
+        # Store original model probabilities (temperature=1, no sampling processors applied)
+        model_ps = logits.softmax(-1)
+
+        # Logit processors (temperature, top-k, top-p). We can chain these!
+        ps = (logits / temperature).softmax(-1)
+        ps = _apply_top_k(ps, model)
+        ps = _apply_top_p(ps, model)
+
+        # Sample the next token and gather the log probabilities
+        if do_sample:  # Sampling
+            next_token_ids = (
+                torch.multinomial(ps, 1) * active_seqs + pad_token_id * ~active_seqs
+            )
+        else:  # Greedy decoding
+            next_token_ids = (
+                torch.argmax(ps, dim=-1).unsqueeze(-1) * active_seqs
+                + pad_token_id * ~active_seqs
+            )
+        next_token_logprobs = ps.gather(-1, next_token_ids).log()
+        next_token_model_logprobs = model_ps.gather(-1, next_token_ids).log()
+
+        input_ids = torch.cat([input_ids, next_token_ids], dim=-1)
+        scores[:, i] = (next_token_logprobs * active_seqs).squeeze()
+        logprobs[:, i] = (next_token_model_logprobs * active_seqs).squeeze()
+
+        active_seqs &= ~torch.isin(next_token_ids, eos_token_ids)
+        if active_seqs.sum() == 0:
+            break
+    return input_ids.detach().cpu(), scores[:, : i + 1], logprobs[:, : i + 1]
+
+
+def generate(model, **kwargs):
+    """
+    Sampling strategy - multinomial sampling with temperature and optional top-k/top-p filtering.
+    Simple implementation with proper KV caching support.
+
+    Args:
+        model: The language model
+        model_kwargs: Model keyword arguments from the tokenizer
+        generation_config: Generation configuration
+        temperature: Sampling temperature (higher = more random)
+        top_k: Only consider top-k most probable tokens
+        top_p: Only consider tokens with cumulative probability <= top_p
+        **kwargs: Additional arguments
+
+    Returns:
+        Generated token IDs
+    """
+    generation_config = model.generation_config
+    max_new_tokens = kwargs.get("max_new_tokens", generation_config.max_new_tokens)
+    max_new_tokens = 512 if max_new_tokens is None else max_new_tokens
+    do_sample = kwargs.get("do_sample", True)
+    eos_token_ids = kwargs.get("eos_token_ids", generation_config.eos_token_id)
+    if eos_token_ids is None:
+        raise ValueError(
+            "Model generation config does not have an EOS token id. You must provide it to generate() with the eos_token_ids argument."
+        )
+    eos_token_ids = torch.as_tensor(eos_token_ids, device=model.device)
+    if eos_token_ids is not None and eos_token_ids.ndim == 0:
+        eos_token_ids = eos_token_ids.unsqueeze(0)
+
+    pad_token_id = kwargs.get(
+        "pad_token_id",
+        generation_config.pad_token_id
+        if generation_config.pad_token_id is not None
+        else eos_token_ids[0],
+    )
+    bos_token_id = kwargs.get("bos_token_id", generation_config.bos_token_id)
+    if bos_token_id is None:
+        raise ValueError(
+            "Model generation config does not have a BOS token id. You must provide it to generate() with the bos_token_id argument."
+        )
+    temperature = kwargs.get("temperature", 1.0)
+    return_dict = kwargs.get("return_dict_in_generate", False)
+
+    generated_ids, scores, logprobs = sampling(
+        model_kwargs=kwargs,
+        model=model,
+        eos_token_ids=eos_token_ids,
+        pad_token_id=pad_token_id,
+        bos_token_id=bos_token_id,
+        do_sample=do_sample,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+    )
+
+    if return_dict:
+        return {
+            "sequences": generated_ids,
+            "scores": scores,
+            "logprobs": logprobs,
+        }
+    else:
+        return generated_ids

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "transformers_version": "4.40.1"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80521b40281d6ce74e35c9282c22539e75aa0ac8578892b2a59955ef78d55da1
+size 269060552

special_tokens_map.json

@@ -0,0 +1,42 @@
+{
+  "additional_special_tokens": [
+    "<|endoftext|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<repo_name>",
+    "<reponame>",
+    "<file_sep>",
+    "<filename>",
+    "<gh_stars>",
+    "<issue_start>",
+    "<issue_comment>",
+    "<issue_closed>",
+    "<jupyter_start>",
+    "<jupyter_text>",
+    "<jupyter_code>",
+    "<jupyter_output>",
+    "<jupyter_script>",
+    "<empty_output>"
+  ],
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,167 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<repo_name>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<reponame>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "<file_sep>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "content": "<filename>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "7": {
+      "content": "<gh_stars>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "8": {
+      "content": "<issue_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "9": {
+      "content": "<issue_comment>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "10": {
+      "content": "<issue_closed>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "11": {
+      "content": "<jupyter_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "12": {
+      "content": "<jupyter_text>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "13": {
+      "content": "<jupyter_code>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "14": {
+      "content": "<jupyter_output>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "15": {
+      "content": "<jupyter_script>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16": {
+      "content": "<empty_output>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|endoftext|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<repo_name>",
+    "<reponame>",
+    "<file_sep>",
+    "<filename>",
+    "<gh_stars>",
+    "<issue_start>",
+    "<issue_comment>",
+    "<issue_closed>",
+    "<jupyter_start>",
+    "<jupyter_text>",
+    "<jupyter_code>",
+    "<jupyter_output>",
+    "<jupyter_script>",
+    "<empty_output>"
+  ],
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 8192,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>",
+  "vocab_size": 49152
+}