gptmed 0.0.1__py3-none-any.whl

gptmed/inference/sampling.py

@@ -0,0 +1,261 @@
+"""
+Sampling Strategies for Text Generation
+
+PURPOSE:
+Different methods to select the next token during generation.
+Each strategy has trade-offs between quality, diversity, and speed.
+
+WHAT THIS FILE DOES:
+1. Greedy sampling: Always pick highest probability (deterministic)
+2. Temperature sampling: Control randomness
+3. Top-k sampling: Sample from top k tokens only
+4. Top-p (nucleus) sampling: Sample from cumulative probability p
+
+WHY DIFFERENT STRATEGIES:
+- Greedy: Fast, deterministic, but boring and repetitive
+- Temperature: Simple randomness control
+- Top-k: Prevents sampling very unlikely tokens
+- Top-p: More adaptive than top-k (adjusts to probability distribution)
+
+PACKAGES USED:
+- torch: PyTorch tensors and operations
+
+FILES FROM THIS PROJECT:
+- None (utility functions)
+
+COMMON ISSUES:
+- Temperature too low → boring, repetitive
+- Temperature too high → incoherent
+- Top-k too small → limited diversity
+- Top-p too low → truncates good options
+"""
+
+import torch
+import torch.nn.functional as F
+
+
+def greedy_sample(logits: torch.Tensor) -> torch.Tensor:
+    """
+    Greedy sampling: Always pick the highest probability token.
+
+    Args:
+        logits: Logits from model [batch_size, vocab_size]
+
+    Returns:
+        Next token IDs [batch_size]
+
+    Pros: Fast, deterministic, reproducible
+    Cons: Boring, repetitive, gets stuck in loops
+
+    Use when: You want deterministic outputs or testing
+    """
+    return torch.argmax(logits, dim=-1)
+
+
+def temperature_sample(logits: torch.Tensor, temperature: float) -> torch.Tensor:
+    """
+    Temperature sampling: Scale logits before softmax.
+
+    Args:
+        logits: Logits from model [batch_size, vocab_size]
+        temperature: Temperature parameter (>0)
+
+    Returns:
+        Sampled token IDs [batch_size]
+
+    How it works:
+    - temperature = 1.0: No change (normal sampling)
+    - temperature < 1.0: More conservative (peaks sharper)
+    - temperature > 1.0: More random (distribution flatter)
+
+    Example:
+    - temp=0.1: Almost greedy
+    - temp=0.7: Balanced (recommended)
+    - temp=1.5: Very creative
+
+    Pros: Simple, interpretable
+    Cons: No control over tail probabilities
+    """
+    if temperature == 0.0:
+        return greedy_sample(logits)
+
+    # Scale logits by temperature
+    scaled_logits = logits / temperature
+
+    # Convert to probabilities
+    probs = F.softmax(scaled_logits, dim=-1)
+
+    # Sample from distribution
+    next_token = torch.multinomial(probs, num_samples=1).squeeze(-1)
+
+    return next_token
+
+
+def top_k_sample(logits: torch.Tensor, k: int, temperature: float = 1.0) -> torch.Tensor:
+    """
+    Top-k sampling: Only sample from top k tokens.
+
+    Args:
+        logits: Logits from model [batch_size, vocab_size]
+        k: Number of top tokens to keep
+        temperature: Temperature scaling
+
+    Returns:
+        Sampled token IDs [batch_size]
+
+    How it works:
+    - Keep only top k highest probability tokens
+    - Set all other tokens to -inf (zero probability)
+    - Sample from remaining tokens
+
+    Why it helps:
+    - Prevents sampling very unlikely tokens (noise)
+    - Reduces incoherent outputs
+
+    Typical values:
+    - k=1: Greedy
+    - k=10: Very conservative
+    - k=50: Balanced
+    - k=100: More diverse
+
+    Limitation:
+    - Fixed k doesn't adapt to probability distribution
+    - If top-1 has 99% probability, k=50 wastes options
+    """
+    if k == 0 or k >= logits.size(-1):
+        # No filtering
+        return temperature_sample(logits, temperature)
+
+    # Get top k values and indices
+    top_k_logits, top_k_indices = torch.topk(logits, k, dim=-1)
+
+    # Create filtered logits (set non-top-k to -inf)
+    filtered_logits = torch.full_like(logits, float("-inf"))
+    filtered_logits.scatter_(-1, top_k_indices, top_k_logits)
+
+    # Sample with temperature
+    return temperature_sample(filtered_logits, temperature)
+
+
+def top_p_sample(logits: torch.Tensor, p: float, temperature: float = 1.0) -> torch.Tensor:
+    """
+    Top-p (nucleus) sampling: Sample from smallest set with cumulative prob >= p.
+
+    Args:
+        logits: Logits from model [batch_size, vocab_size]
+        p: Cumulative probability threshold (0 < p <= 1)
+        temperature: Temperature scaling
+
+    Returns:
+        Sampled token IDs [batch_size]
+
+    How it works:
+    1. Sort tokens by probability (descending)
+    2. Find smallest set where cumulative probability >= p
+    3. Sample only from this set
+
+    Why better than top-k:
+    - Adapts to probability distribution
+    - When model is confident (one token has 90%), nucleus is small
+    - When uncertain, nucleus is larger (more options)
+
+    Typical values:
+    - p=0.9: Conservative (90% probability mass)
+    - p=0.95: Balanced (recommended)
+    - p=0.99: More diverse
+
+    Used in: GPT-3, ChatGPT, most modern LLMs
+    """
+    if p >= 1.0:
+        # No filtering
+        return temperature_sample(logits, temperature)
+
+    # Scale by temperature first
+    scaled_logits = logits / temperature
+
+    # Convert to probabilities
+    probs = F.softmax(scaled_logits, dim=-1)
+
+    # Sort probabilities descending
+    sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
+
+    # Compute cumulative probabilities
+    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+
+    # Find cutoff: first position where cumsum > p
+    # Shift right by 1 to keep at least one token
+    sorted_indices_to_remove = cumulative_probs > p
+    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+    sorted_indices_to_remove[..., 0] = False
+
+    # Create mask in original order
+    indices_to_remove = sorted_indices_to_remove.scatter(
+        -1, sorted_indices, sorted_indices_to_remove
+    )
+
+    # Set removed indices to -inf
+    filtered_logits = scaled_logits.clone()
+    filtered_logits[indices_to_remove] = float("-inf")
+
+    # Sample from filtered distribution
+    probs = F.softmax(filtered_logits, dim=-1)
+    next_token = torch.multinomial(probs, num_samples=1).squeeze(-1)
+
+    return next_token
+
+
+def sample_next_token(
+    logits: torch.Tensor, temperature: float = 1.0, top_k: int = 0, top_p: float = 1.0
+) -> torch.Tensor:
+    """
+    Unified sampling function combining temperature, top-k, and top-p.
+
+    Args:
+        logits: Logits from model [batch_size, vocab_size]
+        temperature: Temperature parameter
+        top_k: Top-k filtering (0 = disabled)
+        top_p: Top-p filtering (1.0 = disabled)
+
+    Returns:
+        Sampled token IDs [batch_size]
+
+    Order of operations:
+    1. Temperature scaling
+    2. Top-k filtering (if enabled)
+    3. Top-p filtering (if enabled)
+    4. Sample from remaining distribution
+    """
+    # Greedy if temperature is 0
+    if temperature == 0.0:
+        return greedy_sample(logits)
+
+    # Apply temperature
+    scaled_logits = logits / temperature
+
+    # Apply top-k if enabled
+    if top_k > 0 and top_k < logits.size(-1):
+        top_k_logits, top_k_indices = torch.topk(scaled_logits, top_k, dim=-1)
+        filtered_logits = torch.full_like(scaled_logits, float("-inf"))
+        filtered_logits.scatter_(-1, top_k_indices, top_k_logits)
+        scaled_logits = filtered_logits
+
+    # Apply top-p if enabled
+    if top_p < 1.0:
+        probs = F.softmax(scaled_logits, dim=-1)
+        sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
+        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] = False
+
+        indices_to_remove = sorted_indices_to_remove.scatter(
+            -1, sorted_indices, sorted_indices_to_remove
+        )
+        scaled_logits[indices_to_remove] = float("-inf")
+
+    # Sample
+    probs = F.softmax(scaled_logits, dim=-1)
+    next_token = torch.multinomial(probs, num_samples=1).squeeze(-1)
+
+    return next_token

gptmed/model/__init__.py

@@ -0,0 +1,9 @@
+"""
+MedLLM Model Package
+
+This package contains the GPT-based transformer architecture for medical QA.
+"""
+
+from llm_med.model.architecture import GPTTransformer
+
+__all__ = ["GPTTransformer"]

gptmed/model/architecture/__init__.py

@@ -0,0 +1,35 @@
+"""
+Model Architecture Components - Package Initializer
+
+PURPOSE:
+Makes it easy to import transformer components from a single location.
+
+WHAT THIS FILE DOES:
+Exposes main classes for import:
+- from model.architecture import GPTTransformer
+- from model.architecture import TransformerDecoderBlock
+- from model.architecture import MultiHeadAttention
+
+PACKAGES USED:
+- None (just Python imports)
+
+FILES FROM THIS PROJECT:
+- All components in this architecture/ directory
+"""
+
+from .embeddings import TokenEmbedding, PositionalEmbedding, TokenPositionalEmbedding
+from .attention import MultiHeadAttention, create_causal_mask
+from .feedforward import FeedForward
+from .decoder_block import TransformerDecoderBlock
+from .transformer import GPTTransformer
+
+__all__ = [
+    "TokenEmbedding",
+    "PositionalEmbedding",
+    "TokenPositionalEmbedding",
+    "MultiHeadAttention",
+    "create_causal_mask",
+    "FeedForward",
+    "TransformerDecoderBlock",
+    "GPTTransformer",
+]

gptmed/model/architecture/attention.py

@@ -0,0 +1,188 @@
+"""
+Multi-Head Causal Self-Attention
+
+PURPOSE:
+This is the core mechanism that allows the model to attend to different parts
+of the input sequence. "Causal" means the model can only look at previous tokens,
+not future ones (essential for next-token prediction).
+
+WHAT THIS STEP DOES:
+1. Linear projections: Create Query, Key, Value matrices
+   - Input: [batch_size, seq_len, d_model]
+   - Q, K, V each: [batch_size, seq_len, d_model]
+
+2. Split into multiple heads
+   - Reshape to: [batch_size, n_heads, seq_len, d_head]
+   - where d_head = d_model / n_heads
+
+3. Scaled dot-product attention
+   - Compute attention scores: Q @ K^T / sqrt(d_head)
+   - Apply causal mask (CRITICAL: prevents looking at future)
+   - Softmax to get attention weights
+   - Apply to values: attention_weights @ V
+
+4. Concatenate heads and project back
+   - Output: [batch_size, seq_len, d_model]
+
+PACKAGES USED:
+- torch: PyTorch tensors and operations
+- torch.nn: Linear layers, Dropout
+- torch.nn.functional: Softmax
+- math: sqrt for scaling
+
+FILES FROM THIS PROJECT:
+- None (this is a base component)
+
+TENSOR SHAPES EXPLAINED:
+- n_heads: Number of attention heads (4-8)
+- d_head: Dimension per head (d_model / n_heads)
+- Causal mask: Lower triangular matrix [seq_len, seq_len]
+
+COMMON FAILURE MODES TO AVOID:
+- Missing causal mask → model cheats by seeing future tokens
+- Wrong mask shape → silent failures or crashes
+- Not scaling attention scores → vanishing/exploding gradients
+- Forgetting dropout → overfitting
+- Wrong tensor transpose/reshape → incorrect attention patterns
+- Not masking padding tokens → attending to meaningless tokens
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+
+class MultiHeadAttention(nn.Module):
+    """
+    Multi-Head Causal Self-Attention.
+
+    This is the CORE of the transformer. Understanding this is critical.
+
+    Key concept: Attention lets each token "look at" other tokens to gather context.
+    "Causal" means token at position i can ONLY look at positions <= i (not future).
+
+    Tensor shape flow:
+        Input:  [batch_size, seq_len, d_model]
+        Q,K,V:  [batch_size, seq_len, d_model]
+        Split:  [batch_size, n_heads, seq_len, d_head]
+        Attention: [batch_size, n_heads, seq_len, seq_len]
+        Output: [batch_size, seq_len, d_model]
+    """
+
+    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
+        """
+        Args:
+            d_model: Model dimension
+            n_heads: Number of attention heads
+            dropout: Dropout probability
+        """
+        super().__init__()
+
+        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
+
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.d_head = d_model // n_heads  # Dimension per head
+
+        # Linear projections for Q, K, V
+        # We use separate projections for each, not combined
+        self.q_linear = nn.Linear(d_model, d_model)
+        self.k_linear = nn.Linear(d_model, d_model)
+        self.v_linear = nn.Linear(d_model, d_model)
+
+        # Output projection
+        self.out_linear = nn.Linear(d_model, d_model)
+
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+
+        # Scaling factor for attention scores
+        self.scale = math.sqrt(self.d_head)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
+        """
+        Args:
+            x: Input tensor [batch_size, seq_len, d_model]
+            mask: Causal mask [seq_len, seq_len] or [batch_size, seq_len, seq_len]
+
+        Returns:
+            Output tensor [batch_size, seq_len, d_model]
+        """
+        batch_size, seq_len, d_model = x.size()
+
+        # 1. Linear projections
+        # Each: [batch_size, seq_len, d_model]
+        Q = self.q_linear(x)
+        K = self.k_linear(x)
+        V = self.v_linear(x)
+
+        # 2. Split into multiple heads
+        # Reshape: [batch_size, seq_len, n_heads, d_head]
+        # Then transpose: [batch_size, n_heads, seq_len, d_head]
+        Q = Q.view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)
+        K = K.view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)
+        V = V.view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2)
+
+        # 3. Scaled dot-product attention
+        # Q @ K^T: [batch_size, n_heads, seq_len, seq_len]
+        attention_scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
+
+        # 4. Apply causal mask (CRITICAL!)
+        # This prevents position i from attending to positions > i
+        if mask is not None:
+            # mask should be [seq_len, seq_len] or [batch_size, 1, seq_len, seq_len]
+            attention_scores = attention_scores.masked_fill(mask == 0, float("-inf"))
+
+        # 5. Softmax to get attention weights
+        # [batch_size, n_heads, seq_len, seq_len]
+        attention_weights = F.softmax(attention_scores, dim=-1)
+
+        # Apply dropout to attention weights (standard practice)
+        attention_weights = self.dropout(attention_weights)
+
+        # 6. Apply attention to values
+        # [batch_size, n_heads, seq_len, d_head]
+        attended_values = torch.matmul(attention_weights, V)
+
+        # 7. Concatenate heads
+        # Transpose back: [batch_size, seq_len, n_heads, d_head]
+        # Then reshape: [batch_size, seq_len, d_model]
+        attended_values = attended_values.transpose(1, 2).contiguous()
+        attended_values = attended_values.view(batch_size, seq_len, self.d_model)
+
+        # 8. Final linear projection
+        output = self.out_linear(attended_values)
+
+        return output
+
+
+def create_causal_mask(seq_len: int, device: torch.device = None) -> torch.Tensor:
+    """
+    Create a causal mask for autoregressive generation.
+
+    The mask is a lower triangular matrix:
+        [[1, 0, 0, 0],
+         [1, 1, 0, 0],
+         [1, 1, 1, 0],
+         [1, 1, 1, 1]]
+
+    This ensures position i can only attend to positions <= i.
+
+    Args:
+        seq_len: Sequence length
+        device: Device to create tensor on
+
+    Returns:
+        Causal mask [seq_len, seq_len]
+
+    Why this works:
+    - Position 0 can see position 0 only
+    - Position 1 can see positions 0, 1
+    - Position 2 can see positions 0, 1, 2
+    - etc.
+
+    This is the ESSENCE of causal language modeling!
+    """
+    mask = torch.tril(torch.ones(seq_len, seq_len, device=device))
+    return mask  # [seq_len, seq_len]

gptmed/model/architecture/decoder_block.py

@@ -0,0 +1,130 @@
+"""
+Transformer Decoder Block
+
+PURPOSE:
+Combines attention, feed-forward network, residual connections, and layer
+normalization into a single reusable block. Multiple blocks are stacked
+to form the full transformer.
+
+WHAT THIS STEP DOES:
+1. Multi-head causal self-attention
+   - With residual connection: x + attention(x)
+   - With layer normalization
+
+2. Feed-forward network
+   - With residual connection: x + ffn(x)
+   - With layer normalization
+
+Architecture pattern (Pre-LN vs Post-LN):
+We use Pre-LN (normalize before sublayer) because it's more stable:
+  x = x + attention(LayerNorm(x))
+  x = x + ffn(LayerNorm(x))
+
+PACKAGES USED:
+- torch: PyTorch tensors
+- torch.nn: Module, LayerNorm, Dropout
+
+FILES FROM THIS PROJECT:
+- architecture/attention.py: Multi-head attention module
+- architecture/feedforward.py: FFN module
+
+TENSOR SHAPES:
+- Input: [batch_size, seq_len, d_model]
+- Output: [batch_size, seq_len, d_model] (unchanged)
+
+COMMON FAILURE MODES TO AVOID:
+- Post-LN instead of Pre-LN → training instability
+- Forgetting residual connections → vanishing gradients
+- Wrong LayerNorm dimension → incorrect normalization
+- Dropout too high → underfitting
+- Dropout too low → overfitting
+"""
+
+import torch
+import torch.nn as nn
+
+from .attention import MultiHeadAttention
+from .feedforward import FeedForward
+
+
+class TransformerDecoderBlock(nn.Module):
+    """
+    Single Transformer Decoder Block.
+
+    This is one "layer" of the transformer. We stack multiple of these.
+
+    Architecture (Pre-LN):
+        1. LayerNorm → Multi-Head Attention → Residual
+        2. LayerNorm → Feed-Forward → Residual
+
+    Tensor shape flow:
+        Input:  [batch_size, seq_len, d_model]
+        Output: [batch_size, seq_len, d_model] (same shape)
+
+    Why Pre-LN instead of Post-LN?
+    - Pre-LN: Normalize BEFORE sublayer → more stable gradients
+    - Post-LN: Normalize AFTER sublayer → can have training instability
+    - GPT-2 and modern transformers use Pre-LN
+    """
+
+    def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.1):
+        """
+        Args:
+            d_model: Model dimension
+            n_heads: Number of attention heads
+            d_ff: Feed-forward hidden dimension
+            dropout: Dropout probability
+        """
+        super().__init__()
+
+        # Multi-head self-attention
+        self.attention = MultiHeadAttention(d_model, n_heads, dropout)
+
+        # Feed-forward network
+        self.feed_forward = FeedForward(d_model, d_ff, dropout)
+
+        # Layer normalization (one for each sublayer)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+
+        # Dropout for residual connections
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
+        """
+        Args:
+            x: Input tensor [batch_size, seq_len, d_model]
+            mask: Causal mask [seq_len, seq_len]
+
+        Returns:
+            Output tensor [batch_size, seq_len, d_model]
+
+        Step-by-step explanation:
+        1. Normalize input
+        2. Apply attention
+        3. Add residual (skip connection)
+        4. Normalize result
+        5. Apply feed-forward
+        6. Add residual (skip connection)
+        """
+        # Sublayer 1: Self-Attention with residual
+        # Pre-LN: normalize first
+        normed = self.norm1(x)
+
+        # Apply attention
+        attention_output = self.attention(normed, mask)
+
+        # Residual connection: x + attention(norm(x))
+        x = x + self.dropout(attention_output)
+
+        # Sublayer 2: Feed-Forward with residual
+        # Pre-LN: normalize first
+        normed = self.norm2(x)
+
+        # Apply feed-forward
+        ff_output = self.feed_forward(normed)
+
+        # Residual connection: x + ffn(norm(x))
+        x = x + self.dropout(ff_output)
+
+        return x