gptmed 0.0.1__py3-none-any.whl

gptmed/model/architecture/embeddings.py

@@ -0,0 +1,146 @@
+"""
+Token and Positional Embeddings
+
+PURPOSE:
+This file converts token IDs (integers) into continuous vector representations
+that the Transformer can process. It also adds positional information so the
+model knows the order of tokens (since attention has no inherent notion of position).
+
+WHAT THIS STEP DOES:
+1. Token Embedding: Maps each token ID to a learned vector of size d_model
+   - Input: [batch_size, seq_len] of token IDs
+   - Output: [batch_size, seq_len, d_model] of vectors
+
+2. Positional Embedding: Adds position information to each token
+   - Two approaches: learned embeddings or sinusoidal encoding
+   - Same shape as token embeddings: [batch_size, seq_len, d_model]
+
+3. Combines both: token_emb + pos_emb
+   - Final output: [batch_size, seq_len, d_model]
+
+PACKAGES USED:
+- torch: PyTorch tensors and neural network modules
+- torch.nn: Embedding layer, Dropout
+
+FILES FROM THIS PROJECT:
+- None (this is a base component)
+
+TENSOR SHAPES EXPLAINED:
+- batch_size: Number of sequences processed together
+- seq_len: Length of each sequence (512 in our case)
+- vocab_size: Size of tokenizer vocabulary (8000)
+- d_model: Embedding dimension (256-512)
+
+COMMON FAILURE MODES TO AVOID:
+- Forgetting dropout → overfitting
+- Wrong positional encoding dimension → shape mismatch
+- Not scaling embeddings → training instability
+- Using fixed positions > max_seq_len → index out of bounds
+"""
+
+import torch
+import torch.nn as nn
+import math
+
+
+class TokenEmbedding(nn.Module):
+    """
+    Convert token IDs to embeddings.
+
+    Tensor shape flow:
+        Input:  [batch_size, seq_len] (token IDs)
+        Output: [batch_size, seq_len, d_model] (embeddings)
+    """
+
+    def __init__(self, vocab_size: int, d_model: int):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        self.d_model = d_model
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Token IDs [batch_size, seq_len]
+
+        Returns:
+            Embeddings [batch_size, seq_len, d_model]
+
+        Note: We scale by sqrt(d_model) following the original Transformer paper.
+        This prevents embeddings from being too small relative to positional encodings.
+        """
+        return self.embedding(x) * math.sqrt(self.d_model)
+
+
+class PositionalEmbedding(nn.Module):
+    """
+    Learned positional embeddings.
+
+    Alternative to sinusoidal encoding - the model learns optimal position representations.
+    GPT uses learned embeddings, so we follow that.
+
+    Tensor shape flow:
+        Input:  [batch_size, seq_len]
+        Output: [batch_size, seq_len, d_model]
+    """
+
+    def __init__(self, max_seq_len: int, d_model: int):
+        super().__init__()
+        self.embedding = nn.Embedding(max_seq_len, d_model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Token IDs [batch_size, seq_len] (we only use seq_len from this)
+
+        Returns:
+            Position embeddings [batch_size, seq_len, d_model]
+        """
+        batch_size, seq_len = x.size()
+
+        # Create position indices: [0, 1, 2, ..., seq_len-1]
+        positions = torch.arange(seq_len, device=x.device).unsqueeze(0)  # [1, seq_len]
+
+        # Expand to batch: [batch_size, seq_len]
+        positions = positions.expand(batch_size, seq_len)
+
+        return self.embedding(positions)
+
+
+class TokenPositionalEmbedding(nn.Module):
+    """
+    Combined token + positional embeddings.
+
+    This is what the transformer actually uses.
+
+    Tensor shape flow:
+        Input:  [batch_size, seq_len] (token IDs)
+        Output: [batch_size, seq_len, d_model] (combined embeddings)
+    """
+
+    def __init__(self, vocab_size: int, d_model: int, max_seq_len: int, dropout: float = 0.1):
+        super().__init__()
+        self.token_embedding = TokenEmbedding(vocab_size, d_model)
+        self.positional_embedding = PositionalEmbedding(max_seq_len, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Token IDs [batch_size, seq_len]
+
+        Returns:
+            Combined embeddings [batch_size, seq_len, d_model]
+        """
+        # Get token embeddings: [batch_size, seq_len, d_model]
+        token_emb = self.token_embedding(x)
+
+        # Get positional embeddings: [batch_size, seq_len, d_model]
+        pos_emb = self.positional_embedding(x)
+
+        # Combine by addition
+        # Why addition? Because we want the model to learn relationships
+        # between both "what" (token) and "where" (position)
+        embeddings = token_emb + pos_emb
+
+        # Apply dropout for regularization
+        return self.dropout(embeddings)

gptmed/model/architecture/feedforward.py

@@ -0,0 +1,109 @@
+"""
+Feed-Forward Network (FFN)
+
+PURPOSE:
+Applies position-wise transformations to each token independently.
+This adds non-linear processing power to the model beyond what attention provides.
+
+WHAT THIS STEP DOES:
+1. First linear projection: Expand dimension
+   - Input: [batch_size, seq_len, d_model]
+   - Output: [batch_size, seq_len, d_ff]
+   - Typically d_ff = 4 * d_model (expansion)
+
+2. Non-linear activation (GELU or ReLU)
+   - Introduces non-linearity
+   - GELU is smoother than ReLU, often better for transformers
+
+3. Second linear projection: Project back
+   - Input: [batch_size, seq_len, d_ff]
+   - Output: [batch_size, seq_len, d_model]
+
+4. Dropout for regularization
+
+PACKAGES USED:
+- torch: PyTorch tensors
+- torch.nn: Linear, Dropout, GELU/ReLU
+
+FILES FROM THIS PROJECT:
+- None (this is a base component)
+
+TENSOR SHAPES EXPLAINED:
+- d_ff: Hidden dimension in FFN (usually 4 * d_model)
+- For d_model=256, d_ff=1024
+- For d_model=512, d_ff=2048
+
+COMMON FAILURE MODES TO AVOID:
+- d_ff too small → insufficient expressiveness
+- d_ff too large → OOM on GPU, overfitting
+- Forgetting activation → just linear transformation (useless)
+- Using ReLU with high learning rate → dead neurons
+"""
+
+import torch
+import torch.nn as nn
+
+
+class FeedForward(nn.Module):
+    """
+    Position-wise Feed-Forward Network.
+
+    Applies the same transformation to each position independently.
+    This is why it's called "position-wise" - no interaction between positions.
+
+    Architecture:
+        Linear(d_model → d_ff) → GELU → Linear(d_ff → d_model) → Dropout
+
+    Tensor shape flow:
+        Input:  [batch_size, seq_len, d_model]
+        Hidden: [batch_size, seq_len, d_ff]
+        Output: [batch_size, seq_len, d_model]
+    """
+
+    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
+        """
+        Args:
+            d_model: Model dimension (embedding size)
+            d_ff: Feed-forward hidden dimension (typically 4 * d_model)
+            dropout: Dropout probability
+        """
+        super().__init__()
+
+        # First linear projection: expand
+        self.linear1 = nn.Linear(d_model, d_ff)
+
+        # Non-linear activation
+        # GELU (Gaussian Error Linear Unit) - smoother than ReLU
+        # Used in GPT-2, BERT, and most modern transformers
+        self.activation = nn.GELU()
+
+        # Second linear projection: compress back
+        self.linear2 = nn.Linear(d_ff, d_model)
+
+        # Dropout for regularization
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Input tensor [batch_size, seq_len, d_model]
+
+        Returns:
+            Output tensor [batch_size, seq_len, d_model]
+
+        Note: Each position is processed identically and independently.
+        No attention or interaction between different positions here.
+        """
+        # Expand: [batch_size, seq_len, d_model] → [batch_size, seq_len, d_ff]
+        x = self.linear1(x)
+
+        # Non-linearity
+        x = self.activation(x)
+
+        # Compress back: [batch_size, seq_len, d_ff] → [batch_size, seq_len, d_model]
+        x = self.linear2(x)
+
+        # Dropout
+        x = self.dropout(x)
+
+        return x

gptmed/model/architecture/transformer.py

@@ -0,0 +1,204 @@
+"""
+Full GPT-Style Transformer Model
+
+PURPOSE:
+Assembles all components (embeddings, decoder blocks, output projection)
+into a complete causal language model for next-token prediction.
+
+WHAT THIS STEP DOES:
+1. Token + Positional Embeddings
+   - Convert input IDs to vectors: [batch_size, seq_len, d_model]
+
+2. Stack of N Decoder Blocks
+   - Each block applies self-attention + FFN
+   - Typically N = 4-6 layers for our hardware
+
+3. Final Layer Normalization
+   - Stabilize outputs before projection
+
+4. Output Projection (Language Modeling Head)
+   - Project to vocabulary: [batch_size, seq_len, vocab_size]
+   - No activation (raw logits for CrossEntropyLoss)
+
+5. Forward Pass
+   - Input: token IDs [batch_size, seq_len]
+   - Output: logits [batch_size, seq_len, vocab_size]
+
+PACKAGES USED:
+- torch: PyTorch tensors
+- torch.nn: Module, Linear, LayerNorm, ModuleList
+
+FILES FROM THIS PROJECT:
+- architecture/embeddings.py: TokenPositionalEmbedding
+- architecture/decoder_block.py: TransformerDecoderBlock
+- configs/model_config.py: Hyperparameters (d_model, n_layers, etc.)
+
+TENSOR SHAPES:
+- Input IDs: [batch_size, seq_len] (integers)
+- Embeddings: [batch_size, seq_len, d_model]
+- After blocks: [batch_size, seq_len, d_model]
+- Logits: [batch_size, seq_len, vocab_size]
+
+HYPERPARAMETERS (for GTX 1080):
+- vocab_size: 8000 (from tokenizer)
+- d_model: 256-512 (embedding dimension)
+- n_layers: 4-6 (number of transformer blocks)
+- n_heads: 4-8 (attention heads)
+- d_ff: 4 * d_model (FFN hidden size)
+- dropout: 0.1-0.2
+- max_seq_len: 512
+
+COMMON FAILURE MODES TO AVOID:
+- Not tying embeddings and output weights → slower convergence
+- Too many layers → OOM or slow training
+- d_model not divisible by n_heads → shape mismatch
+- Missing final LayerNorm → unstable outputs
+- Forgetting to handle padding mask → attending to padding
+"""
+
+import torch
+import torch.nn as nn
+
+from .embeddings import TokenPositionalEmbedding
+from .decoder_block import TransformerDecoderBlock
+from .attention import create_causal_mask
+
+
+class GPTTransformer(nn.Module):
+    """
+    GPT-style Causal Language Model.
+
+    This is the COMPLETE model. Everything comes together here.
+
+    Architecture:
+        1. Token + Positional Embeddings
+        2. N x Decoder Blocks
+        3. Final LayerNorm
+        4. LM Head (projects to vocab)
+
+    Training objective: Next-token prediction
+        - Given tokens [t0, t1, t2, ..., tn]
+        - Predict [t1, t2, t3, ..., tn+1]
+
+    Tensor flow:
+        Input:  [batch_size, seq_len] token IDs
+        Embed:  [batch_size, seq_len, d_model]
+        Blocks: [batch_size, seq_len, d_model]
+        Logits: [batch_size, seq_len, vocab_size]
+    """
+
+    def __init__(self, config):
+        """
+        Args:
+            config: ModelConfig object with hyperparameters
+        """
+        super().__init__()
+
+        self.config = config
+
+        # 1. Embeddings (token + positional)
+        self.embeddings = TokenPositionalEmbedding(
+            vocab_size=config.vocab_size,
+            d_model=config.d_model,
+            max_seq_len=config.max_seq_len,
+            dropout=config.dropout,
+        )
+
+        # 2. Stack of decoder blocks
+        self.decoder_blocks = nn.ModuleList(
+            [
+                TransformerDecoderBlock(
+                    d_model=config.d_model,
+                    n_heads=config.n_heads,
+                    d_ff=config.d_ff,
+                    dropout=config.dropout,
+                )
+                for _ in range(config.n_layers)
+            ]
+        )
+
+        # 3. Final layer normalization
+        self.final_norm = nn.LayerNorm(config.d_model)
+
+        # 4. Language modeling head (output projection)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+
+        # Weight tying: Share weights between embeddings and output projection
+        # This is standard practice in language models
+        # Why? Reduces parameters and improves generalization
+        self.lm_head.weight = self.embeddings.token_embedding.embedding.weight
+
+        # Initialize weights
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        """
+        Initialize weights using Xavier/Glorot initialization.
+
+        This is CRITICAL for training stability.
+        Poor initialization → vanishing/exploding gradients
+        """
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            torch.nn.init.ones_(module.weight)
+            torch.nn.init.zeros_(module.bias)
+
+    def forward(self, input_ids: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
+        """
+        Forward pass through the model.
+
+        Args:
+            input_ids: Token IDs [batch_size, seq_len]
+            mask: Optional causal mask [seq_len, seq_len]
+
+        Returns:
+            Logits [batch_size, seq_len, vocab_size]
+
+        Step-by-step:
+        1. Convert token IDs to embeddings
+        2. Pass through each decoder block
+        3. Apply final normalization
+        4. Project to vocabulary
+        """
+        batch_size, seq_len = input_ids.size()
+
+        # Create causal mask if not provided
+        if mask is None:
+            mask = create_causal_mask(seq_len, device=input_ids.device)
+
+        # 1. Embeddings: [batch_size, seq_len] → [batch_size, seq_len, d_model]
+        x = self.embeddings(input_ids)
+
+        # 2. Pass through decoder blocks
+        for block in self.decoder_blocks:
+            x = block(x, mask)
+
+        # 3. Final layer norm
+        x = self.final_norm(x)
+
+        # 4. Project to vocabulary: [batch_size, seq_len, d_model] → [batch_size, seq_len, vocab_size]
+        logits = self.lm_head(x)
+
+        return logits
+
+    def count_parameters(self) -> int:
+        """Count total trainable parameters."""
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+    def get_num_params(self, non_embedding: bool = False) -> int:
+        """
+        Get parameter count.
+
+        Args:
+            non_embedding: If True, exclude embedding parameters
+        """
+        n_params = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n_params -= self.embeddings.token_embedding.embedding.weight.numel()
+            n_params -= self.embeddings.positional_embedding.embedding.weight.numel()
+        return n_params

gptmed/model/configs/__init__.py

@@ -0,0 +1,17 @@
+"""
+Model configuration module
+"""
+
+from llm_med.model.configs.model_config import (
+    ModelConfig,
+    get_tiny_config,
+    get_small_config,
+    get_medium_config,
+)
+
+__all__ = [
+    "ModelConfig",
+    "get_tiny_config",
+    "get_small_config",
+    "get_medium_config",
+]

gptmed/model/configs/model_config.py

@@ -0,0 +1,155 @@
+"""
+Model Configuration
+
+PURPOSE:
+Central place to store all model hyperparameters. Makes it easy to experiment
+with different model sizes without changing code.
+
+WHAT THIS FILE CONTAINS:
+1. ModelConfig dataclass with all hyperparameters:
+   - vocab_size: From tokenizer (8000)
+   - d_model: Embedding/hidden dimension
+   - n_layers: Number of transformer blocks
+   - n_heads: Number of attention heads
+   - d_ff: Feed-forward hidden dimension
+   - dropout: Dropout probability
+   - max_seq_len: Maximum sequence length
+
+2. Predefined configurations:
+   - Tiny: For quick testing (d_model=128, n_layers=2)
+   - Small: For GTX 1080 training (d_model=256, n_layers=4)
+   - Medium: Larger if memory allows (d_model=512, n_layers=6)
+
+PACKAGES USED:
+- dataclasses: For clean config structure
+- json: For saving/loading configs
+
+FILES FROM THIS PROJECT:
+- None (this defines configs for other files to use)
+
+DESIGN DECISIONS:
+- d_model must be divisible by n_heads
+- d_ff typically 4 * d_model (expansion ratio)
+- dropout 0.1-0.2 (too high → underfitting, too low → overfitting)
+- max_seq_len matches tokenization (512)
+
+MEMORY ESTIMATION (approximate):
+- Model parameters ≈ 12 * n_layers * d_model^2
+- Small config: ~10M parameters (~40MB)
+- Medium config: ~40M parameters (~160MB)
+- Fits comfortably in 8GB VRAM with batch_size=16-32
+"""
+
+from dataclasses import dataclass
+import json
+from pathlib import Path
+
+
+@dataclass
+class ModelConfig:
+    """
+    Transformer model configuration.
+
+    All hyperparameters in one place for easy experimentation.
+    """
+
+    # Vocabulary
+    vocab_size: int = 8000  # From SentencePiece tokenizer
+
+    # Architecture
+    d_model: int = 256  # Embedding/hidden dimension
+    n_layers: int = 4  # Number of transformer blocks
+    n_heads: int = 4  # Number of attention heads
+    d_ff: int = 1024  # FFN hidden dimension (4 * d_model)
+
+    # Regularization
+    dropout: float = 0.1  # Dropout probability
+
+    # Sequence
+    max_seq_len: int = 512  # Maximum sequence length
+
+    # Special tokens (from tokenizer)
+    pad_token_id: int = 0
+    unk_token_id: int = 1
+    bos_token_id: int = 2
+    eos_token_id: int = 3
+
+    def __post_init__(self):
+        """Validate configuration."""
+        assert (
+            self.d_model % self.n_heads == 0
+        ), f"d_model ({self.d_model}) must be divisible by n_heads ({self.n_heads})"
+        assert self.d_model > 0, "d_model must be positive"
+        assert self.n_layers > 0, "n_layers must be positive"
+        assert 0.0 <= self.dropout < 1.0, "dropout must be in [0, 1)"
+
+    @property
+    def d_head(self) -> int:
+        """Dimension per attention head."""
+        return self.d_model // self.n_heads
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary."""
+        return {
+            "vocab_size": self.vocab_size,
+            "d_model": self.d_model,
+            "n_layers": self.n_layers,
+            "n_heads": self.n_heads,
+            "d_ff": self.d_ff,
+            "dropout": self.dropout,
+            "max_seq_len": self.max_seq_len,
+            "pad_token_id": self.pad_token_id,
+            "unk_token_id": self.unk_token_id,
+            "bos_token_id": self.bos_token_id,
+            "eos_token_id": self.eos_token_id,
+        }
+
+    def save(self, path: Path):
+        """Save config to JSON file."""
+        with open(path, "w") as f:
+            json.dump(self.to_dict(), f, indent=2)
+
+    @classmethod
+    def from_dict(cls, config_dict: dict) -> "ModelConfig":
+        """Load from dictionary."""
+        return cls(**config_dict)
+
+    @classmethod
+    def from_file(cls, path: Path) -> "ModelConfig":
+        """Load from JSON file."""
+        with open(path, "r") as f:
+            config_dict = json.load(f)
+        return cls.from_dict(config_dict)
+
+
+# Predefined configurations for different use cases
+
+
+def get_tiny_config() -> ModelConfig:
+    """
+    Tiny model for quick testing and debugging.
+    ~2M parameters, very fast training.
+    """
+    return ModelConfig(
+        vocab_size=8000, d_model=128, n_layers=2, n_heads=4, d_ff=512, dropout=0.1, max_seq_len=512
+    )
+
+
+def get_small_config() -> ModelConfig:
+    """
+    Small model for GTX 1080 (8GB VRAM).
+    ~10M parameters, good balance of speed and capacity.
+    """
+    return ModelConfig(
+        vocab_size=8000, d_model=256, n_layers=4, n_heads=4, d_ff=1024, dropout=0.1, max_seq_len=512
+    )
+
+
+def get_medium_config() -> ModelConfig:
+    """
+    Medium model for GPUs with more memory.
+    ~40M parameters, better quality but slower.
+    """
+    return ModelConfig(
+        vocab_size=8000, d_model=512, n_layers=6, n_heads=8, d_ff=2048, dropout=0.1, max_seq_len=512
+    )

gptmed/tokenizer/__init__.py

@@ -0,0 +1,7 @@
+"""
+Tokenizer module for MedLLM
+"""
+
+from llm_med.tokenizer.train_tokenizer import train_sentencepiece_tokenizer
+
+__all__ = ["train_sentencepiece_tokenizer"]