gptmed 0.0.1__py3-none-any.whl

gptmed/inference/decoding_utils.py

@@ -0,0 +1,190 @@
+"""
+Decoding Utilities
+
+PURPOSE:
+Helper functions for text generation:
+- Repetition penalty
+- N-gram blocking
+- Stopping criteria
+
+WHAT THIS FILE DOES:
+1. Apply repetition penalty to discourage repeated tokens
+2. Block n-gram repetition (prevents "the the the")
+3. Check stopping conditions
+
+WHY THESE ARE NEEDED:
+- Models often get stuck in repetition loops
+- "The patient has has has has..."
+- N-gram blocking prevents this
+- Repetition penalty makes it less likely
+
+PACKAGES USED:
+- torch: PyTorch tensors
+
+FILES FROM THIS PROJECT:
+- None (utility functions)
+"""
+
+import torch
+from typing import List, Set
+
+
+def apply_repetition_penalty(
+    logits: torch.Tensor, generated_tokens: torch.Tensor, penalty: float = 1.0
+) -> torch.Tensor:
+    """
+    Apply repetition penalty to discourage repeated tokens.
+
+    Args:
+        logits: Current logits [batch_size, vocab_size]
+        generated_tokens: Previously generated tokens [batch_size, seq_len]
+        penalty: Penalty factor (>1.0 penalizes, <1.0 encourages)
+
+    Returns:
+        Modified logits
+
+    How it works:
+    - For each token that appeared before:
+      - If its logit is positive: divide by penalty
+      - If its logit is negative: multiply by penalty
+    - This makes repeated tokens less likely
+
+    Typical values:
+    - penalty=1.0: No penalty
+    - penalty=1.1: Mild penalty
+    - penalty=1.2: Moderate (recommended)
+    - penalty=1.5: Strong penalty
+
+    Warning: Too high penalty can make model avoid common words!
+    """
+    if penalty == 1.0:
+        return logits
+
+    batch_size = logits.size(0)
+
+    for batch_idx in range(batch_size):
+        # Get unique tokens in this sequence
+        unique_tokens = generated_tokens[batch_idx].unique()
+
+        for token_id in unique_tokens:
+            token_id = token_id.item()
+
+            # Apply penalty
+            if logits[batch_idx, token_id] > 0:
+                logits[batch_idx, token_id] /= penalty
+            else:
+                logits[batch_idx, token_id] *= penalty
+
+    return logits
+
+
+def get_ngrams(tokens: List[int], n: int) -> Set[tuple]:
+    """
+    Extract all n-grams from token list.
+
+    Args:
+        tokens: List of token IDs
+        n: N-gram size
+
+    Returns:
+        Set of n-grams (tuples)
+
+    Example:
+        tokens = [1, 2, 3, 4]
+        get_ngrams(tokens, 2) = {(1,2), (2,3), (3,4)}
+    """
+    if len(tokens) < n:
+        return set()
+
+    ngrams = set()
+    for i in range(len(tokens) - n + 1):
+        ngram = tuple(tokens[i : i + n])
+        ngrams.add(ngram)
+
+    return ngrams
+
+
+def block_ngram_repeats(
+    logits: torch.Tensor, generated_tokens: List[int], ngram_size: int = 3
+) -> torch.Tensor:
+    """
+    Block tokens that would create repeated n-grams.
+
+    Args:
+        logits: Current logits [vocab_size]
+        generated_tokens: Previously generated tokens (list)
+        ngram_size: N-gram size to block
+
+    Returns:
+        Modified logits with blocked tokens set to -inf
+
+    How it works:
+    - Look at last (n-1) tokens
+    - Find all tokens that appeared after this (n-1)-gram before
+    - Set their logits to -inf (can't be sampled)
+
+    Example with n=3:
+    - Generated: "The patient has the patient"
+    - Last 2 tokens: "the patient"
+    - Previously after "the patient": "has"
+    - Block "has" from being generated again
+
+    This prevents: "The patient has the patient has the patient has..."
+    """
+    if ngram_size == 0 or len(generated_tokens) < ngram_size:
+        return logits
+
+    # Get context (last n-1 tokens)
+    context = tuple(generated_tokens[-(ngram_size - 1) :])
+
+    # Find all tokens that appeared after this context
+    blocked_tokens = set()
+
+    for i in range(len(generated_tokens) - ngram_size + 1):
+        # Check if this position matches our context
+        if tuple(generated_tokens[i : i + ngram_size - 1]) == context:
+            # The next token creates a repeated n-gram
+            next_token = generated_tokens[i + ngram_size - 1]
+            blocked_tokens.add(next_token)
+
+    # Block these tokens
+    for token_id in blocked_tokens:
+        logits[token_id] = float("-inf")
+
+    return logits
+
+
+def should_stop_generation(
+    generated_tokens: List[int], stop_tokens: List[int], max_length: int, min_length: int
+) -> bool:
+    """
+    Check if generation should stop.
+
+    Args:
+        generated_tokens: Generated token IDs
+        stop_tokens: Token IDs that trigger stopping (e.g., EOS)
+        max_length: Maximum allowed length
+        min_length: Minimum required length
+
+    Returns:
+        True if should stop, False otherwise
+
+    Stopping criteria:
+    1. Reached max_length
+    2. Generated a stop token (and past min_length)
+    """
+    current_length = len(generated_tokens)
+
+    # Must generate at least min_length tokens
+    if current_length < min_length:
+        return False
+
+    # Stop if reached max length
+    if current_length >= max_length:
+        return True
+
+    # Stop if generated a stop token
+    if generated_tokens[-1] in stop_tokens:
+        return True
+
+    return False

gptmed/inference/generation_config.py

@@ -0,0 +1,83 @@
+"""
+Generation Configuration
+
+PURPOSE:
+Hyperparameters for text generation (inference).
+Controls how creative vs conservative the model's outputs are.
+
+WHAT THIS FILE CONTAINS:
+- Temperature: Randomness control
+- Top-k, top-p: Sampling constraints
+- Repetition penalty: Prevent repetitive text
+- Max length: Stop generation
+
+PACKAGES USED:
+- dataclasses: Clean config structure
+
+FILES FROM THIS PROJECT:
+- None (base config)
+
+KEY PARAMETERS EXPLAINED:
+- temperature: 0.0 = greedy, 0.7 = balanced, 1.5 = very creative
+- top_k: Only sample from top k tokens (50-100 typical)
+- top_p: Sample from tokens with cumulative prob p (0.9-0.95 typical)
+- repetition_penalty: >1.0 discourages repetition (1.2 is good)
+"""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class GenerationConfig:
+    """Configuration for text generation."""
+
+    # Sampling strategy
+    temperature: float = 0.8  # Higher = more random (0.0 = greedy)
+    top_k: int = 50  # Only sample from top k tokens (0 = disabled)
+    top_p: float = 0.95  # Nucleus sampling threshold (1.0 = disabled)
+
+    # Repetition control
+    repetition_penalty: float = 1.2  # >1.0 penalizes repetition
+    no_repeat_ngram_size: int = 3  # Block repeating n-grams (0 = disabled)
+
+    # Length control
+    max_length: int = 200  # Maximum tokens to generate
+    min_length: int = 10  # Minimum tokens to generate
+
+    # Stopping criteria
+    stop_tokens: list = None  # Token IDs that stop generation
+
+    # Special tokens
+    bos_token_id: int = 2  # Beginning of sequence
+    eos_token_id: int = 3  # End of sequence
+    pad_token_id: int = 0  # Padding
+
+    def __post_init__(self):
+        """Validate config."""
+        if self.stop_tokens is None:
+            self.stop_tokens = [self.eos_token_id]
+
+        assert self.temperature >= 0.0, "temperature must be >= 0"
+        assert self.top_k >= 0, "top_k must be >= 0"
+        assert 0.0 <= self.top_p <= 1.0, "top_p must be in [0, 1]"
+        assert self.repetition_penalty >= 1.0, "repetition_penalty must be >= 1.0"
+
+
+def get_greedy_config() -> GenerationConfig:
+    """Greedy decoding (deterministic, picks highest prob)."""
+    return GenerationConfig(temperature=0.0, top_k=0, top_p=1.0, repetition_penalty=1.0)
+
+
+def get_balanced_config() -> GenerationConfig:
+    """Balanced sampling (good default)."""
+    return GenerationConfig(temperature=0.8, top_k=50, top_p=0.95, repetition_penalty=1.2)
+
+
+def get_creative_config() -> GenerationConfig:
+    """Creative sampling (more diverse, less coherent)."""
+    return GenerationConfig(temperature=1.2, top_k=100, top_p=0.95, repetition_penalty=1.3)
+
+
+def get_conservative_config() -> GenerationConfig:
+    """Conservative sampling (safe, coherent, less diverse)."""
+    return GenerationConfig(temperature=0.5, top_k=30, top_p=0.9, repetition_penalty=1.1)

gptmed/inference/generator.py

@@ -0,0 +1,253 @@
+"""
+Text Generator
+
+PURPOSE:
+Main class for text generation using trained GPT model.
+Combines model loading, tokenization, and decoding strategies.
+
+WHAT THIS FILE DOES:
+1. Load trained model from checkpoint
+2. Generate text autoregressively (token by token)
+3. Apply sampling strategies and repetition control
+4. Convert tokens back to text
+
+GENERATION PROCESS:
+1. Start with prompt tokens
+2. Loop:
+   a. Model forward pass → logits
+   b. Apply repetition penalty
+   c. Sample next token
+   d. Append to sequence
+   e. Check stopping criteria
+3. Decode tokens to text
+
+PACKAGES USED:
+- torch: PyTorch
+- sentencepiece: Tokenizer
+
+FILES FROM THIS PROJECT:
+- model/architecture/transformer.py: GPT model
+- model/configs/model_config.py: Model config
+- inference/sampling.py: Sampling strategies
+- inference/decoding_utils.py: Repetition control
+- utils/checkpoints.py: Load checkpoints
+"""
+
+import torch
+import sentencepiece as spm
+from pathlib import Path
+from typing import List, Optional
+
+from llm_med.model.architecture import GPTTransformer
+from llm_med.model.configs.model_config import ModelConfig
+from llm_med.inference.generation_config import GenerationConfig
+from llm_med.inference.sampling import sample_next_token
+from llm_med.inference.decoding_utils import (
+    apply_repetition_penalty,
+    block_ngram_repeats,
+    should_stop_generation,
+)
+
+
+class TextGenerator:
+    """
+    Text generation with trained GPT model.
+
+    This is your interface for inference.
+    """
+
+    def __init__(
+        self, model: GPTTransformer, tokenizer: spm.SentencePieceProcessor, device: str = "cuda"
+    ):
+        """
+        Args:
+            model: Trained GPT model
+            tokenizer: SentencePiece tokenizer
+            device: Device to run on
+        """
+        self.model = model.to(device)
+        self.model.eval()  # Set to evaluation mode
+        self.tokenizer = tokenizer
+        self.device = device
+
+    @classmethod
+    def from_checkpoint(cls, checkpoint_path: Path, tokenizer_path: Path, device: str = "cuda"):
+        """
+        Load generator from checkpoint.
+
+        Args:
+            checkpoint_path: Path to model checkpoint
+            tokenizer_path: Path to tokenizer .model file
+            device: Device to load on
+
+        Returns:
+            TextGenerator instance
+        """
+        # Load checkpoint
+        print(f"Loading checkpoint: {checkpoint_path}")
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+
+        # Create model from config
+        model_config_dict = checkpoint["model_config"]
+        model_config = ModelConfig(**model_config_dict)
+        model = GPTTransformer(model_config)
+
+        # Load weights
+        model.load_state_dict(checkpoint["model_state_dict"])
+        model.to(device)
+        model.eval()
+
+        print(f"Model loaded (step {checkpoint['step']})")
+        print(f"Validation loss: {checkpoint.get('val_loss', 'N/A')}")
+
+        # Load tokenizer
+        print(f"Loading tokenizer: {tokenizer_path}")
+        tokenizer = spm.SentencePieceProcessor()
+        tokenizer.load(str(tokenizer_path))
+
+        return cls(model, tokenizer, device)
+
+    def encode(self, text: str) -> List[int]:
+        """Encode text to token IDs."""
+        return self.tokenizer.encode_as_ids(text)
+
+    def decode(self, token_ids: List[int]) -> str:
+        """Decode token IDs to text."""
+        return self.tokenizer.decode_ids(token_ids)
+
+    def generate(
+        self, prompt: str, gen_config: GenerationConfig = None, verbose: bool = False
+    ) -> str:
+        """
+        Generate text from a prompt.
+
+        Args:
+            prompt: Input text prompt
+            gen_config: Generation configuration
+            verbose: Print generation progress
+
+        Returns:
+            Generated text
+
+        Process:
+        1. Encode prompt to tokens
+        2. Generate tokens autoregressively
+        3. Decode back to text
+        """
+        if gen_config is None:
+            gen_config = GenerationConfig()
+
+        # Encode prompt
+        prompt_tokens = self.encode(prompt)
+
+        if verbose:
+            print(f"Prompt: {prompt}")
+            print(f"Prompt tokens: {prompt_tokens}")
+            print(f"Generating...")
+
+        # Generate tokens
+        generated_tokens = self.generate_tokens(
+            prompt_tokens=prompt_tokens, gen_config=gen_config, verbose=verbose
+        )
+
+        # Decode to text
+        generated_text = self.decode(generated_tokens)
+
+        if verbose:
+            print(f"\nGenerated {len(generated_tokens)} tokens")
+            print(f"Output: {generated_text}")
+
+        return generated_text
+
+    def generate_tokens(
+        self, prompt_tokens: List[int], gen_config: GenerationConfig, verbose: bool = False
+    ) -> List[int]:
+        """
+        Generate tokens autoregressively.
+
+        Args:
+            prompt_tokens: Input token IDs
+            gen_config: Generation config
+            verbose: Print progress
+
+        Returns:
+            List of generated token IDs (including prompt)
+        """
+        # Start with prompt
+        generated = prompt_tokens.copy()
+
+        # Generation loop
+        with torch.no_grad():
+            for step in range(gen_config.max_length):
+                # Check stopping criteria
+                if should_stop_generation(
+                    generated_tokens=generated,
+                    stop_tokens=gen_config.stop_tokens,
+                    max_length=gen_config.max_length,
+                    min_length=gen_config.min_length,
+                ):
+                    break
+
+                # Prepare input (last max_seq_len tokens)
+                max_seq_len = self.model.config.max_seq_len
+                input_ids = generated[-max_seq_len:]
+                input_tensor = torch.tensor([input_ids], device=self.device)
+
+                # Forward pass
+                logits = self.model(input_tensor)
+
+                # Get logits for last position
+                next_token_logits = logits[0, -1, :]  # [vocab_size]
+
+                # Apply repetition penalty
+                if gen_config.repetition_penalty != 1.0:
+                    generated_tensor = torch.tensor([generated], device=self.device)
+                    next_token_logits = apply_repetition_penalty(
+                        next_token_logits.unsqueeze(0),
+                        generated_tensor,
+                        gen_config.repetition_penalty,
+                    ).squeeze(0)
+
+                # Block n-gram repeats
+                if gen_config.no_repeat_ngram_size > 0:
+                    next_token_logits = block_ngram_repeats(
+                        next_token_logits, generated, gen_config.no_repeat_ngram_size
+                    )
+
+                # Sample next token
+                next_token = sample_next_token(
+                    next_token_logits.unsqueeze(0),
+                    temperature=gen_config.temperature,
+                    top_k=gen_config.top_k,
+                    top_p=gen_config.top_p,
+                )
+
+                next_token_id = next_token.item()
+                generated.append(next_token_id)
+
+                if verbose and step % 10 == 0:
+                    partial_text = self.decode(generated)
+                    print(f"Step {step}: {partial_text}")
+
+        return generated
+
+    def generate_batch(self, prompts: List[str], gen_config: GenerationConfig = None) -> List[str]:
+        """
+        Generate text for multiple prompts.
+
+        Args:
+            prompts: List of input prompts
+            gen_config: Generation config
+
+        Returns:
+            List of generated texts
+        """
+        if gen_config is None:
+            gen_config = GenerationConfig()
+
+        results = []
+        for prompt in prompts:
+            output = self.generate(prompt, gen_config, verbose=False)
+            results.append(output)
+
+        return results