Upload generation.py with huggingface_hub

generation.py

@@ -0,0 +1,205 @@
+"""
+Text Generation Utilities for ASA Models
+
+Simple, dependency-free text generation with common decoding strategies.
+
+Repository: https://github.com/DigitalDaimyo/AddressedStateAttention
+"""
+
+import torch
+import torch.nn.functional as F
+from typing import Optional, Set, Tuple, List
+
+
+__all__ = ['generate']
+
+
+def _forward_logits(model, input_ids, attention_mask=None):
+    """Extract logits from various model output formats."""
+    out = model(input_ids, attention_mask=attention_mask) if attention_mask is not None else model(input_ids)
+    
+    if isinstance(out, torch.Tensor):
+        return out
+    if isinstance(out, (tuple, list)):
+        return out[0]
+    if isinstance(out, dict):
+        for key in ["logits", "out", "y", "pred"]:
+            if key in out:
+                return out[key]
+    raise TypeError(f"Unrecognized model output type: {type(out)}")
+
+
+def _apply_repetition_penalty(logits: torch.Tensor, input_ids: torch.Tensor, penalty: float):
+    """Apply repetition penalty to logits (GPT-2 style)."""
+    if penalty is None or penalty == 1.0:
+        return logits
+    
+    B = logits.size(0)
+    for b in range(B):
+        prev_tokens = torch.unique(input_ids[b])
+        l = logits[b, prev_tokens]
+        logits[b, prev_tokens] = torch.where(l < 0, l * penalty, l / penalty)
+    return logits
+
+
+def _top_k_top_p_filtering(
+    logits: torch.Tensor, 
+    top_k: int = 0, 
+    top_p: float = 1.0, 
+    min_tokens_to_keep: int = 1
+):
+    """Filter logits using top-k and nucleus (top-p) filtering."""
+    B, V = logits.shape
+    top_k = int(top_k) if top_k is not None else 0
+    top_p = float(top_p) if top_p is not None else 1.0
+
+    if top_k > 0 and top_k < V:
+        kth = torch.topk(logits, top_k, dim=-1).values[:, -1].unsqueeze(-1)
+        logits = logits.masked_fill(logits < kth, float("-inf"))
+
+    if top_p < 1.0:
+        sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
+        probs = F.softmax(sorted_logits, dim=-1)
+        cum = probs.cumsum(dim=-1)
+
+        remove = cum > top_p
+        if min_tokens_to_keep > 1:
+            remove[:, :min_tokens_to_keep] = False
+        remove = torch.cat([
+            torch.zeros((B, 1), device=logits.device, dtype=torch.bool), 
+            remove[:, :-1]
+        ], dim=-1)
+
+        sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
+        logits = torch.full_like(logits, float("-inf"))
+        logits.scatter_(dim=-1, index=sorted_idx, src=sorted_logits)
+
+    return logits
+
+
+def _update_seen_ngrams(seen: Set, tokens: List[int], n: int):
+    """Add n-gram to seen set."""
+    if n > 0 and len(tokens) >= n:
+        seen.add(tuple(tokens[-n:]))
+
+
+def _seed_seen_ngrams(input_ids: torch.Tensor, n: int) -> Set:
+    """Initialize seen n-grams from input."""
+    seen = set()
+    if n <= 0:
+        return seen
+    tokens = input_ids[0].tolist()
+    if len(tokens) >= n:
+        for i in range(len(tokens) - n + 1):
+            seen.add(tuple(tokens[i:i+n]))
+    return seen
+
+
+def _banned_from_seen(seen: Set, input_ids: torch.Tensor, n: int) -> Set:
+    """Get tokens banned by n-gram constraint."""
+    if n <= 0 or input_ids.shape[1] < n - 1:
+        return set()
+    
+    prefix = tuple(input_ids[0, -(n - 1):].tolist())
+    banned = set()
+    for ng in seen:
+        if ng[:-1] == prefix:
+            banned.add(ng[-1])
+    return banned
+
+
+@torch.no_grad()
+def generate(
+    model,
+    tokenizer,
+    prompt: str,
+    max_new_tokens: int = 120,
+    max_seq_len: int = 1024,
+    strategy: str = "sample",
+    temperature: float = 1.0,
+    top_k: int = 0,
+    top_p: float = 0.9,
+    repetition_penalty: float = 1.0,
+    no_repeat_ngram_size: int = 0,
+    eos_token_id: Optional[int] = None,
+    device: str = "cuda",
+) -> str:
+    """
+    Generate text from a prompt using various decoding strategies.
+    
+    Args:
+        model: ASA language model
+        tokenizer: HuggingFace tokenizer
+        prompt: Input text prompt
+        max_new_tokens: Maximum tokens to generate
+        max_seq_len: Maximum sequence length (truncates context if exceeded)
+        strategy: "greedy" or "sample"
+        temperature: Sampling temperature (higher = more random)
+        top_k: Keep only top k tokens (0 = disabled)
+        top_p: Nucleus sampling threshold (1.0 = disabled)
+        repetition_penalty: Penalty for repeating tokens (1.0 = disabled)
+        no_repeat_ngram_size: Block repeating n-grams (0 = disabled)
+        eos_token_id: Stop generation at this token
+        device: Device to run on
+        
+    Returns:
+        Generated text (including prompt)
+        
+    Example:
+        >>> text = generate(
+        ...     model, tokenizer,
+        ...     prompt="The capital of France is",
+        ...     max_new_tokens=20,
+        ...     strategy="greedy"
+        ... )
+    """
+    model.eval()
+    
+    enc = tokenizer(prompt, return_tensors="pt")
+    input_ids = enc.input_ids.to(device)
+    
+    if eos_token_id is None:
+        eos_token_id = tokenizer.eos_token_id
+    
+    seen = _seed_seen_ngrams(input_ids, no_repeat_ngram_size)
+    
+    for _ in range(max_new_tokens):
+        # Truncate if exceeding context length
+        if input_ids.shape[1] > max_seq_len:
+            input_ids = input_ids[:, -max_seq_len:]
+            seen = _seed_seen_ngrams(input_ids, no_repeat_ngram_size)
+        
+        logits = _forward_logits(model, input_ids)
+        next_logits = logits[:, -1, :].to(torch.float32).clone()
+        
+        # Apply repetition penalty
+        next_logits = _apply_repetition_penalty(next_logits, input_ids, repetition_penalty)
+        
+        # Block repeated n-grams
+        banned = _banned_from_seen(seen, input_ids, no_repeat_ngram_size)
+        if banned:
+            next_logits[0, list(banned)] = float("-inf")
+        
+        # Decode strategy
+        if strategy == "greedy":
+            next_token = torch.argmax(next_logits, dim=-1, keepdim=True)
+        elif strategy == "sample":
+            temp = max(1e-6, float(temperature))
+            next_logits = next_logits / temp
+            next_logits = _top_k_top_p_filtering(next_logits, top_k=top_k, top_p=top_p)
+            probs = F.softmax(next_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+        else:
+            raise ValueError(f"Unknown strategy '{strategy}'. Use 'greedy' or 'sample'.")
+        
+        input_ids = torch.cat([input_ids, next_token], dim=1)
+        
+        # Update n-gram tracking
+        tokens = input_ids[0].tolist()
+        _update_seen_ngrams(seen, tokens, no_repeat_ngram_size)
+        
+        # Check for EOS
+        if eos_token_id is not None and next_token.item() == eos_token_id:
+            break
+    
+    return tokenizer.decode(input_ids[0], skip_special_tokens=False)