gptmed 0.0.1__py3-none-any.whl

gptmed/tokenizer/tokenize_data.py

@@ -0,0 +1,286 @@
+"""
+Dataset Tokenization
+
+Applies trained tokenizer to processed text and creates train/validation splits.
+
+Design decisions:
+- Sequence length: 512 tokens (fits GTX 1080, captures most Q&A pairs)
+- Train/val split: 90/10 (enough validation data to detect overfitting)
+- Padding: Left-pad or truncate (causal LM sees left-to-right)
+- No data augmentation: Keep it simple for Phase 1
+
+Common failure modes:
+- Truncating answers → model never learns to generate long responses
+- Too short sequences → can't learn context
+- Too long sequences → OOM on GPU
+- Wrong padding → breaks attention masks
+"""
+
+import argparse
+import json
+from pathlib import Path
+from typing import List, Tuple
+import random
+
+import sentencepiece as spm
+import numpy as np
+
+
+def load_tokenizer(model_path: Path) -> spm.SentencePieceProcessor:
+    """Load trained SentencePiece tokenizer."""
+    sp = spm.SentencePieceProcessor()
+    sp.load(str(model_path))
+    return sp
+
+
+def tokenize_text(text: str, tokenizer: spm.SentencePieceProcessor) -> List[int]:
+    """
+    Tokenize text to IDs.
+
+    Args:
+        text: Input text
+        tokenizer: SentencePiece processor
+
+    Returns:
+        List of token IDs
+    """
+    return tokenizer.encode_as_ids(text)
+
+
+def create_sequences(token_ids: List[int], max_length: int, stride: int = None) -> List[List[int]]:
+    """
+    Split token IDs into fixed-length sequences.
+
+    Args:
+        token_ids: Full token ID sequence
+        max_length: Maximum sequence length
+        stride: Stride for overlapping windows (None = no overlap)
+
+    Returns:
+        List of sequences
+
+    Design note: For causal LM, we create non-overlapping chunks.
+    Each chunk is a separate training example for next-token prediction.
+    """
+    if stride is None:
+        stride = max_length  # No overlap
+
+    sequences = []
+    for i in range(0, len(token_ids) - max_length + 1, stride):
+        seq = token_ids[i : i + max_length]
+        sequences.append(seq)
+
+    # Handle remaining tokens (last incomplete sequence)
+    remainder = len(token_ids) % stride
+    if remainder > 0 and len(sequences) > 0:
+        # Include last incomplete sequence if it's at least half the max_length
+        last_start = len(token_ids) - remainder
+        if remainder >= max_length // 2:
+            last_seq = token_ids[last_start:]
+            # Pad to max_length
+            last_seq = last_seq + [0] * (max_length - len(last_seq))
+            sequences.append(last_seq)
+
+    return sequences
+
+
+def analyze_lengths(text_file: Path, tokenizer: spm.SentencePieceProcessor) -> dict:
+    """
+    Analyze token length distribution to choose optimal sequence length.
+
+    This helps avoid:
+    - Truncating too much data (information loss)
+    - Wasting memory on padding (inefficiency)
+    """
+    lengths = []
+
+    with open(text_file, "r", encoding="utf-8") as f:
+        content = f.read()
+
+    # Split by double newline (our document separator)
+    documents = content.split("\n\n")
+
+    for doc in documents:
+        if doc.strip():
+            tokens = tokenizer.encode_as_ids(doc.strip())
+            lengths.append(len(tokens))
+
+    lengths = np.array(lengths)
+
+    stats = {
+        "num_documents": len(lengths),
+        "mean": float(np.mean(lengths)),
+        "median": float(np.median(lengths)),
+        "std": float(np.std(lengths)),
+        "min": int(np.min(lengths)),
+        "max": int(np.max(lengths)),
+        "percentile_50": float(np.percentile(lengths, 50)),
+        "percentile_75": float(np.percentile(lengths, 75)),
+        "percentile_90": float(np.percentile(lengths, 90)),
+        "percentile_95": float(np.percentile(lengths, 95)),
+        "percentile_99": float(np.percentile(lengths, 99)),
+    }
+
+    return stats
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Tokenize MedQuAD dataset with trained tokenizer")
+    parser.add_argument(
+        "--input-file",
+        type=str,
+        default="./data/processed/medquad_simple.txt",
+        help="Input text file",
+    )
+    parser.add_argument(
+        "--tokenizer-model",
+        type=str,
+        default="./tokenizer/medquad_tokenizer.model",
+        help="Path to trained tokenizer model",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="./data/tokenized",
+        help="Output directory for tokenized data",
+    )
+    parser.add_argument(
+        "--max-length",
+        type=int,
+        default=512,
+        help="Maximum sequence length (default: 512 for GTX 1080)",
+    )
+    parser.add_argument(
+        "--train-ratio", type=float, default=0.9, help="Train/val split ratio (default: 0.9)"
+    )
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility")
+
+    args = parser.parse_args()
+
+    # Set random seed
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+
+    print("=" * 60)
+    print("Dataset Tokenization")
+    print("=" * 60)
+    print(f"Input: {args.input_file}")
+    print(f"Tokenizer: {args.tokenizer_model}")
+    print(f"Max length: {args.max_length}")
+    print(f"Train ratio: {args.train_ratio}")
+    print()
+
+    # Load tokenizer
+    input_file = Path(args.input_file)
+    tokenizer_model = Path(args.tokenizer_model)
+
+    if not input_file.exists():
+        print(f"❌ Error: Input file not found: {input_file}")
+        return
+
+    if not tokenizer_model.exists():
+        print(f"❌ Error: Tokenizer not found: {tokenizer_model}")
+        print("\nPlease train tokenizer first:")
+        print("  python tokenizer/train_tokenizer.py")
+        return
+
+    print("Loading tokenizer...")
+    tokenizer = load_tokenizer(tokenizer_model)
+    print(f"✅ Tokenizer loaded (vocab size: {tokenizer.vocab_size()})")
+
+    # Analyze sequence lengths
+    print("\nAnalyzing sequence lengths...")
+    stats = analyze_lengths(input_file, tokenizer)
+
+    print("\nSequence Length Statistics:")
+    print(f"  Documents: {stats['num_documents']}")
+    print(f"  Mean: {stats['mean']:.1f} tokens")
+    print(f"  Median: {stats['median']:.1f} tokens")
+    print(f"  Std dev: {stats['std']:.1f} tokens")
+    print(f"  Min: {stats['min']} tokens")
+    print(f"  Max: {stats['max']} tokens")
+    print(f"\nPercentiles:")
+    print(f"  50th: {stats['percentile_50']:.1f} tokens")
+    print(f"  75th: {stats['percentile_75']:.1f} tokens")
+    print(f"  90th: {stats['percentile_90']:.1f} tokens")
+    print(f"  95th: {stats['percentile_95']:.1f} tokens")
+    print(f"  99th: {stats['percentile_99']:.1f} tokens")
+
+    # Warning about truncation
+    truncated_pct = (
+        np.array([l for l in stats.values() if isinstance(l, (int, float))]) > args.max_length
+    ).sum()
+    if stats["percentile_95"] > args.max_length:
+        print(
+            f"\n⚠️  WARNING: max_length={args.max_length} will truncate ~{100-95:.0f}% of sequences"
+        )
+        print(f"   Consider increasing to {int(stats['percentile_95'])} to capture 95% of data")
+
+    # Tokenize full dataset
+    print("\nTokenizing dataset...")
+    with open(input_file, "r", encoding="utf-8") as f:
+        full_text = f.read()
+
+    token_ids = tokenizer.encode_as_ids(full_text)
+    print(f"Total tokens: {len(token_ids):,}")
+
+    # Create sequences
+    print(f"\nCreating sequences (max_length={args.max_length})...")
+    sequences = create_sequences(token_ids, max_length=args.max_length)
+    print(f"Total sequences: {len(sequences):,}")
+
+    # Train/val split
+    random.shuffle(sequences)
+    split_idx = int(len(sequences) * args.train_ratio)
+
+    train_sequences = sequences[:split_idx]
+    val_sequences = sequences[split_idx:]
+
+    print(f"\nSplit:")
+    print(f"  Train: {len(train_sequences):,} sequences")
+    print(f"  Val: {len(val_sequences):,} sequences")
+
+    # Save tokenized data
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    print("\nSaving tokenized data...")
+
+    # Save as numpy arrays (efficient for PyTorch DataLoader)
+    train_array = np.array(train_sequences, dtype=np.int32)
+    val_array = np.array(val_sequences, dtype=np.int32)
+
+    np.save(output_dir / "train.npy", train_array)
+    np.save(output_dir / "val.npy", val_array)
+
+    print(f"✅ Train data saved: {output_dir / 'train.npy'}")
+    print(f"✅ Val data saved: {output_dir / 'val.npy'}")
+
+    # Save metadata
+    metadata = {
+        "vocab_size": tokenizer.vocab_size(),
+        "max_length": args.max_length,
+        "num_train_sequences": len(train_sequences),
+        "num_val_sequences": len(val_sequences),
+        "total_tokens": len(token_ids),
+        "length_stats": stats,
+        "tokenizer_model": str(tokenizer_model),
+        "seed": args.seed,
+    }
+
+    with open(output_dir / "metadata.json", "w") as f:
+        json.dump(metadata, f, indent=2)
+
+    print(f"✅ Metadata saved: {output_dir / 'metadata.json'}")
+
+    print("\n" + "=" * 60)
+    print("✅ Tokenization complete!")
+    print("=" * 60)
+    print("\nNext steps (Phase 2):")
+    print("1. Implement Transformer architecture")
+    print("2. Create PyTorch Dataset and DataLoader")
+    print("3. Define training loop")
+
+
+if __name__ == "__main__":
+    main()

gptmed/tokenizer/train_tokenizer.py

@@ -0,0 +1,218 @@
+"""
+SentencePiece Tokenizer Training
+
+Trains a BPE (Byte-Pair Encoding) tokenizer on the processed MedQuAD text.
+
+Design decisions explained:
+- BPE over WordPiece: Better for medical terminology (handles rare words via subwords)
+- Vocab size: Trade-off between model size and expressiveness
+- Character coverage: 0.9995 captures medical unicode (Greek letters, symbols)
+- No pre-tokenization: Let BPE learn from raw text
+
+Common failure modes to avoid:
+- Too small vocab → repetitive, generic text
+- Too large vocab → overfitting, slow training
+- Missing special tokens → can't mark boundaries
+- Wrong normalization → "COVID-19" vs "covid-19" treated differently
+"""
+
+import argparse
+from pathlib import Path
+import sentencepiece as spm
+
+
+def train_sentencepiece_tokenizer(
+    input_file: Path,
+    output_prefix: Path,
+    vocab_size: int = 8000,
+    model_type: str = "bpe",
+    character_coverage: float = 0.9995,
+    add_special_tokens: bool = True,
+):
+    """
+    Train a SentencePiece tokenizer.
+
+    Args:
+        input_file: Path to training text file
+        output_prefix: Output path prefix (will create .model and .vocab files)
+        vocab_size: Vocabulary size (default 8000 for small dataset)
+        model_type: 'bpe' or 'unigram'
+        character_coverage: Character coverage for unicode (0.9995 = medical terms)
+        add_special_tokens: Whether to add custom special tokens
+
+    Design notes:
+    - vocab_size=8000: Good for 40K examples on GTX 1080
+      * Too small (2K): Poor compression, repetitive outputs
+      * Too large (32K): Overfits, wastes memory on rare terms
+    - BPE: Deterministic, easier to debug than unigram
+    - character_coverage=0.9995: Captures medical unicode (μ, α, β, etc.)
+    """
+
+    # Build training command
+    train_args = [
+        f"--input={input_file}",
+        f"--model_prefix={output_prefix}",
+        f"--vocab_size={vocab_size}",
+        f"--model_type={model_type}",
+        f"--character_coverage={character_coverage}",
+        "--pad_id=0",  # <pad> token at ID 0
+        "--unk_id=1",  # <unk> token at ID 1
+        "--bos_id=2",  # <bos> (beginning of sequence) at ID 2
+        "--eos_id=3",  # <eos> (end of sequence) at ID 3
+        "--pad_piece=[PAD]",
+        "--unk_piece=[UNK]",
+        "--bos_piece=[BOS]",
+        "--eos_piece=[EOS]",
+    ]
+
+    # Add user-defined special tokens if requested
+    if add_special_tokens:
+        # These match our text formatting (Q:, A:)
+        user_defined = "[Q],[A]"
+        train_args.append(f"--user_defined_symbols={user_defined}")
+
+    # Normalization rules (CAREFUL: medical text is case-sensitive)
+    # We use minimal normalization to preserve:
+    # - "COVID-19" vs "covid-19"
+    # - Drug names (proper capitalization matters)
+    # - Abbreviations (BP vs bp)
+    train_args.extend(
+        [
+            "--normalization_rule_name=identity",  # No normalization
+            "--remove_extra_whitespaces=true",  # Clean whitespace
+            "--split_by_unicode_script=true",  # Split CJK if present
+            "--split_by_whitespace=true",
+            "--split_by_number=true",  # Split numbers
+            "--max_sentence_length=4096",  # Long medical answers
+        ]
+    )
+
+    print("=" * 60)
+    print("Training SentencePiece Tokenizer")
+    print("=" * 60)
+    print(f"Input: {input_file}")
+    print(f"Output prefix: {output_prefix}")
+    print(f"Vocab size: {vocab_size}")
+    print(f"Model type: {model_type}")
+    print(f"Character coverage: {character_coverage}")
+    print()
+
+    # Train the tokenizer
+    spm.SentencePieceTrainer.Train(" ".join(train_args))
+
+    print("\n✅ Tokenizer training complete!")
+    print(f"Model saved: {output_prefix}.model")
+    print(f"Vocab saved: {output_prefix}.vocab")
+
+    # Load and inspect the tokenizer
+    sp = spm.SentencePieceProcessor()
+    sp.load(f"{output_prefix}.model")
+
+    print("\n" + "=" * 60)
+    print("Tokenizer Statistics")
+    print("=" * 60)
+    print(f"Vocabulary size: {sp.vocab_size()}")
+    print(f"PAD token: {sp.id_to_piece(0)} (ID: 0)")
+    print(f"UNK token: {sp.id_to_piece(1)} (ID: 1)")
+    print(f"BOS token: {sp.id_to_piece(2)} (ID: 2)")
+    print(f"EOS token: {sp.id_to_piece(3)} (ID: 3)")
+
+    # Test tokenization on medical example
+    print("\n" + "=" * 60)
+    print("Sample Tokenization")
+    print("=" * 60)
+
+    test_texts = [
+        "What is diabetes?",
+        "COVID-19 vaccination side effects",
+        "Hypertension treatment guidelines",
+    ]
+
+    for text in test_texts:
+        tokens = sp.encode_as_pieces(text)
+        ids = sp.encode_as_ids(text)
+        print(f"\nText: {text}")
+        print(f"Tokens: {tokens}")
+        print(f"IDs: {ids}")
+        print(f"Token count: {len(tokens)}")
+
+    # Vocabulary inspection
+    print("\n" + "=" * 60)
+    print("Sample Vocabulary (first 20 tokens)")
+    print("=" * 60)
+    for i in range(min(20, sp.vocab_size())):
+        piece = sp.id_to_piece(i)
+        print(f"ID {i:4d}: {piece}")
+
+    return sp
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train SentencePiece tokenizer for MedQuAD")
+    parser.add_argument(
+        "--input-file",
+        type=str,
+        default="./data/processed/medquad_simple.txt",
+        help="Input text file for training",
+    )
+    parser.add_argument(
+        "--output-dir", type=str, default="./tokenizer", help="Output directory for tokenizer files"
+    )
+    parser.add_argument(
+        "--vocab-size",
+        type=int,
+        default=8000,
+        help="Vocabulary size (default: 8000 for small dataset)",
+    )
+    parser.add_argument(
+        "--model-type",
+        type=str,
+        choices=["bpe", "unigram"],
+        default="bpe",
+        help="Tokenizer algorithm",
+    )
+    parser.add_argument(
+        "--character-coverage",
+        type=float,
+        default=0.9995,
+        help="Character coverage (0.9995 = captures medical unicode)",
+    )
+
+    args = parser.parse_args()
+
+    # Validate input file
+    input_file = Path(args.input_file)
+    if not input_file.exists():
+        print(f"❌ Error: Input file not found: {input_file}")
+        print("\nPlease run preprocessing first:")
+        print("  python preprocess.py")
+        return
+
+    # Create output directory
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Output prefix for .model and .vocab files
+    output_prefix = output_dir / "medquad_tokenizer"
+
+    # Train tokenizer
+    tokenizer = train_sentencepiece_tokenizer(
+        input_file=input_file,
+        output_prefix=output_prefix,
+        vocab_size=args.vocab_size,
+        model_type=args.model_type,
+        character_coverage=args.character_coverage,
+        add_special_tokens=True,
+    )
+
+    print("\n" + "=" * 60)
+    print("✅ Tokenizer training complete!")
+    print("=" * 60)
+    print("\nNext steps:")
+    print("1. Inspect vocabulary for medical terms")
+    print("2. Test on sample medical texts")
+    print("3. Tokenize full dataset: python tokenizer/tokenize_data.py")
+
+
+if __name__ == "__main__":
+    main()

gptmed/training/__init__.py

@@ -0,0 +1 @@
+"""Training package."""

gptmed/training/dataset.py

@@ -0,0 +1,183 @@
+"""
+PyTorch Dataset for Tokenized Data
+
+PURPOSE:
+Load tokenized sequences from numpy arrays and provide them as PyTorch tensors
+for training. Handles batching and shuffling efficiently.
+
+WHAT THIS FILE DOES:
+1. Load tokenized data from .npy files
+2. Create training examples for next-token prediction
+3. Provide batches to DataLoader
+
+TRAINING OBJECTIVE EXPLAINED:
+For causal language modeling, each sequence is both input and target:
+- Input:  [token_0, token_1, token_2, ..., token_n-1]
+- Target: [token_1, token_2, token_3, ..., token_n]
+
+The model learns to predict token_i given tokens [0, 1, ..., i-1].
+
+PACKAGES USED:
+- torch: PyTorch tensors and Dataset
+- numpy: Load .npy files
+
+FILES FROM THIS PROJECT:
+- data/tokenized/train.npy (created by tokenize_data.py)
+- data/tokenized/val.npy
+
+TENSOR SHAPES:
+- Loaded data: [num_sequences, seq_len]
+- Each batch: [batch_size, seq_len]
+- Input: [batch_size, seq_len]
+- Target: [batch_size, seq_len]
+
+COMMON FAILURE MODES:
+- Wrong data path → FileNotFoundError
+- Mismatched shapes → crashes during training
+- Not shuffling train data → poor generalization
+- Loading entire dataset into memory → OOM (for large datasets)
+"""
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
+from pathlib import Path
+
+
+class TokenizedDataset(Dataset):
+    """
+    Dataset for tokenized sequences.
+
+    This is a simple in-memory dataset. For larger datasets, you'd use
+    memory-mapped arrays or streaming.
+
+    Each item returns (input_ids, target_ids) for next-token prediction.
+    """
+
+    def __init__(self, data_path: Path):
+        """
+        Args:
+            data_path: Path to .npy file with tokenized sequences
+        """
+        self.data_path = Path(data_path)
+
+        if not self.data_path.exists():
+            raise FileNotFoundError(f"Data file not found: {data_path}")
+
+        # Load tokenized sequences
+        # Shape: [num_sequences, seq_len]
+        self.data = np.load(self.data_path)
+
+        print(f"Loaded dataset from {data_path}")
+        print(f"  Shape: {self.data.shape}")
+        print(f"  Dtype: {self.data.dtype}")
+        print(f"  Num sequences: {len(self.data)}")
+
+    def __len__(self) -> int:
+        """Number of sequences in dataset."""
+        return len(self.data)
+
+    def __getitem__(self, idx: int) -> tuple:
+        """
+        Get a single training example.
+
+        Args:
+            idx: Sequence index
+
+        Returns:
+            (input_ids, target_ids) tuple
+
+        Training objective:
+            input:  [t0, t1, t2, ..., tn-1]
+            target: [t1, t2, t3, ..., tn]
+
+        The model predicts token at position i using tokens [0, ..., i-1].
+        """
+        sequence = self.data[idx]
+
+        # Convert to torch tensor
+        sequence = torch.from_numpy(sequence).long()
+
+        # For causal LM, input and target are the same sequence, shifted by 1
+        # Input:  [0, 1, 2, 3, ..., n-1]
+        # Target: [1, 2, 3, 4, ..., n]
+        input_ids = sequence[:-1]  # All tokens except last
+        target_ids = sequence[1:]  # All tokens except first
+
+        return input_ids, target_ids
+
+
+def create_dataloaders(
+    train_path: Path, val_path: Path, batch_size: int, num_workers: int = 0
+) -> tuple:
+    """
+    Create train and validation dataloaders.
+
+    Args:
+        train_path: Path to train .npy file
+        val_path: Path to validation .npy file
+        batch_size: Batch size
+        num_workers: Number of dataloader workers (0 = main process)
+
+    Returns:
+        (train_loader, val_loader) tuple
+
+    Design decisions:
+    - Shuffle train data (prevents overfitting to sequence order)
+    - Don't shuffle val data (consistent evaluation)
+    - num_workers=0 on small datasets (overhead not worth it)
+    - drop_last=True to avoid partial batches (can cause issues with batch norm)
+    """
+    # Create datasets
+    train_dataset = TokenizedDataset(train_path)
+    val_dataset = TokenizedDataset(val_path)
+
+    # Create dataloaders
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        shuffle=True,  # Shuffle for training
+        num_workers=num_workers,
+        pin_memory=True,  # Faster GPU transfer
+        drop_last=True,  # Drop incomplete last batch
+    )
+
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=batch_size,
+        shuffle=False,  # Don't shuffle validation
+        num_workers=num_workers,
+        pin_memory=True,
+        drop_last=False,  # Keep all validation data
+    )
+
+    print(f"\nDataLoaders created:")
+    print(f"  Train batches: {len(train_loader)}")
+    print(f"  Val batches: {len(val_loader)}")
+    print(f"  Batch size: {batch_size}")
+
+    return train_loader, val_loader
+
+
+def get_batch_info(batch: tuple) -> dict:
+    """
+    Get information about a batch (for debugging).
+
+    Args:
+        batch: (input_ids, target_ids) tuple
+
+    Returns:
+        Dictionary with batch statistics
+    """
+    input_ids, target_ids = batch
+
+    return {
+        "batch_size": input_ids.size(0),
+        "seq_len": input_ids.size(1),
+        "input_shape": tuple(input_ids.shape),
+        "target_shape": tuple(target_ids.shape),
+        "input_dtype": input_ids.dtype,
+        "target_dtype": target_ids.dtype,
+        "input_device": input_ids.device,
+        "target_device": target_ids.device,
+    }