Upload folder using huggingface_hub

README.md

@@ -0,0 +1,128 @@
+---
+library_name: transformers
+tags:
+  - custom_generate
+  - ancestral_sampling
+---
+
+# Multinomial (Ancestral) Sampling simple implementation
+
+## Description
+A clean, hackable implementation of ancestral sampling (multinomial sampling) with full KV cache support. 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.
+
+## Additional 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)
+
+## Output Type changes
+When `return_dict_in_generate=True`, returns a dictionary with:
+- `sequences`: Generated token IDs
+- `scores`: Log probabilities of sampled tokens (with temperature/sampling modifications)
+- `logps`: Original model log probabilities (T=1, no modifications)
+- `prompt_lens`: Length of input prompts
+- `lens`: Final sequence lengths
+
+## 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/ancestral_sampling", trust_remote_code=True)
+
+# With temperature
+gen_out = model.generate(**inputs, custom_generate="manueldeprada/ancestral_sampling", temperature=0.8, trust_remote_code=True)
+
+# With top-k
+gen_out = model.generate(**inputs, custom_generate="manueldeprada/ancestral_sampling", top_k=50, trust_remote_code=True)
+
+# With top-p (nucleus sampling)
+gen_out = model.generate(**inputs, custom_generate="manueldeprada/ancestral_sampling", top_p=0.9, trust_remote_code=True)
+
+# Greedy decoding (no sampling)
+gen_out = model.generate(**inputs, custom_generate="manueldeprada/ancestral_sampling", do_sample=False, trust_remote_code=True)
+
+# Get detailed output with probabilities
+gen_out = model.generate(
+    **inputs, 
+    custom_generate="manueldeprada/ancestral_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.logps}")
+```
+
+## Algorithm
+1. Initialize KV cache and 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
+   - Update KV cache and track sequence completion
+3. Return generated sequences and probability information
+
+## Helper Functions for Custom Generation
+
+The implementation provides two key helper functions that you can use to build your own generation strategies:
+
+### `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`
+
+### `ps_next(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`
+
+### Example: Custom Generation Loop
+
+```py
+from ancestral_sampling.generate import init_gen, ps_next
+
+def custom_generation(model, model_kwargs, max_new_tokens=20, temperature=1.0):
+    # Initialize generation
+    model_kwargs, input_ids = init_gen(model_kwargs, model, max_new_tokens, bos_token_id)
+    
+    for i in range(max_new_tokens):
+        # Get next token logits
+        model_kwargs, logits = ps_next(model, model_kwargs, input_ids)
+        
+        # Your custom logic here
+        probs = (logits / temperature).softmax(-1)
+        next_token = torch.multinomial(probs, 1)
+        
+        # Append token and continue
+        input_ids = torch.cat([input_ids, next_token], dim=-1)
+        
+        # Add your stopping conditions
+        if next_token.item() == eos_token_id:
+            break
+    
+    return input_ids
+```
+

SmolLM2-135M/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

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,192 @@
+import torch
+from transformers import GenerationConfig
+
+
+def ps_next(model, model_kwargs, input_ids):
+    """
+    Auxiliary function to get the next token probabilities and update the KV cache.
+    
+    Args:
+        model: The language model
+        model_kwargs: Model keyword arguments including KV cache
+        input_ids: Current input token IDs
+        T: Temperature for sampling
+    
+    Returns:
+        Updated model_kwargs, probabilities at temperature T, probabilities at T=1
+    """
+    model_inputs = model.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 = model._update_model_kwargs_for_generation(
+        outputs, model_kwargs, is_encoder_decoder=model.config.is_encoder_decoder
+    )
+    del outputs
+    return model_kwargs, logits
+
+def init_gen(model_kwargs, model, max_new_tokens, bos_token_id):
+    """
+    Auxiliary function to initialize the generation process and prepare the KV cache.
+    
+    Args:
+        model_kwargs: Model keyword arguments
+        model: The language model
+        max_new_tokens: Maximum number of new tokens to generate
+    
+    Returns:
+        Model keyword arguments and input token IDs
+    """
+
+    input_ids, model_input_name, model_kwargs = model._prepare_model_inputs(
+        None, bos_token_id, model_kwargs
+    )
+    
+    batch_size = input_ids.shape[0]
+    model._prepare_cache_for_generation(
+        model.generation_config, model_kwargs, None, batch_size, 
+        max_cache_length=max_new_tokens, device=input_ids.device
+    )
+    
+    # Get initial cache position
+    model_kwargs = model._get_initial_cache_position(input_ids.shape[1], input_ids.device, model_kwargs)
+    return model_kwargs, input_ids
+
+def _apply_top_k_top_p(ps, model):
+    if hasattr(model, 'generation_config') and hasattr(model.generation_config, 'top_k') and model.generation_config.top_k is not None:
+        top_k = model.generation_config.top_k
+        top_k = min(top_k, ps.size(-1))
+        indices_to_remove = ps < torch.topk(ps, top_k)[0][..., -1, None]
+        ps[indices_to_remove] = 0.0
+        ps = ps / ps.sum(dim=-1, keepdim=True)
+    
+    # Apply top-p filtering if specified
+    if hasattr(model, 'generation_config') and hasattr(model.generation_config, 'top_p') and model.generation_config.top_p is not None:
+        top_p = model.generation_config.top_p
+        if top_p < 1.0:
+            sorted_probs, sorted_indices = torch.sort(ps, descending=True)
+            cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+            
+            # Remove tokens with cumulative probability above the threshold
+            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
+            ps = ps / ps.sum(dim=-1, keepdim=True)
+    return ps
+
+def ancestral_sampling(model_kwargs, model, eos_token_ids, pad_token_id, bos_token_id, do_sample=True, max_new_tokens=20, T=1.0):
+    """
+    Ancestral 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, max_prompts_len = input_ids.shape
+    prompts_len = (input_ids != pad_token_id).sum(dim=-1)
+
+    # Keeps track of which sequences are finished and their lengths
+    active_seqs = input_ids.new_ones((batch_size, 1), dtype=torch.bool)
+    lens = torch.full((batch_size,), max_prompts_len, dtype=torch.long, device=input_ids.device)
+    # Modified log probabilities of the sequences
+    scores = torch.zeros((batch_size, max_new_tokens), dtype=torch.float32)
+    # Unfiltered sequence log probabilities (T=1, no sampling modifications)
+    logps = torch.zeros((batch_size, max_new_tokens), dtype=torch.float32)
+
+    for i in range(max_new_tokens):
+        # Get the next token probabilities and update the KV cache
+        model_kwargs, logits = ps_next(model, model_kwargs, input_ids)
+        # Original model probabilities (T=1, no sampling modifications)
+        model_ps = logits.softmax(-1)
+        # Sampling probabilities (T, with sampling modifications)
+        ps = (logits/T).softmax(-1)
+        ps = _apply_top_k_top_p(ps, model)
+        
+        # Sample the next token and gather the log probabilities
+        if do_sample:
+            next_token_ids = torch.multinomial(ps, 1) * active_seqs + pad_token_id * ~active_seqs
+        else:
+            next_token_ids = torch.argmax(ps, dim=-1).unsqueeze(-1) * active_seqs + pad_token_id * ~active_seqs
+        next_token_logps = ps.gather(-1, next_token_ids).log()
+        next_token_model_logps = model_ps.gather(-1, next_token_ids).log()
+        
+        input_ids = torch.cat([input_ids, next_token_ids], dim=-1)
+        scores[:, i] = (next_token_logps * active_seqs).squeeze()
+        logps[:, i] = (next_token_model_logps * active_seqs).squeeze()
+        
+        lens += active_seqs.squeeze(-1).long()
+        active_seqs &= ~torch.isin(next_token_ids, eos_token_ids)
+        if active_seqs.sum() == 0:
+            break    
+    return input_ids.detach().cpu(), scores[:,:i+1], logps[:,:i+1], prompts_len, lens.tolist()
+
+def generate(model, **kwargs):
+    """
+    Ancestral 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)
+    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.")
+    T = kwargs.get('temperature', 1.0)
+    return_dict = kwargs.get('return_dict_in_generate', False)
+
+    generated_ids, scores, logps, prompt_lens, lens = ancestral_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,
+        T=T,
+    )
+
+    if return_dict:
+        return {
+            "sequences": generated_ids,
+            "scores": scores,
+            "logps": logps,
+            "prompt_lens": prompt_lens,
+            "lens": lens,
+        }
+    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
+}