gptmed 0.0.1__py3-none-any.whl

gptmed/__init__.py

@@ -0,0 +1,37 @@
+"""
+llm-med: A lightweight medical question-answering language model
+
+This package provides a GPT-based transformer architecture trained on the MedQuAD dataset
+for medical domain question answering.
+
+Main Components:
+- model: GPT transformer architecture
+- inference: Text generation and sampling
+- training: Training loop and utilities
+- tokenizer: SentencePiece tokenizer
+- configs: Configuration management
+- utils: Utility functions
+
+Example:
+    >>> from llm_med.model.architecture import GPTTransformer
+    >>> from llm_med.model.configs.model_config import get_small_config
+    >>> from llm_med.inference.generator import TextGenerator
+    >>> 
+    >>> config = get_small_config()
+    >>> model = GPTTransformer(config)
+"""
+
+__version__ = "0.2.0"
+__author__ = "Sanjog Sigdel"
+__email__ = "sigdelsanjog@gmail.com"
+
+# Expose main components at package level for convenience
+from llm_med.model.architecture import GPTTransformer
+from llm_med.model.configs.model_config import ModelConfig, get_small_config, get_tiny_config
+
+__all__ = [
+    "GPTTransformer",
+    "ModelConfig",
+    "get_small_config",
+    "get_tiny_config",
+]

gptmed/configs/__init__.py

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

gptmed/configs/train_config.py

@@ -0,0 +1,154 @@
+"""
+Training Configuration
+
+PURPOSE:
+Central place for all training hyperparameters. Separating training config
+from model config follows separation of concerns - you can train the same
+model architecture with different training strategies.
+
+WHAT THIS FILE CONTAINS:
+- Batch size, learning rate, epochs
+- Optimizer settings (weight decay, betas)
+- Learning rate schedule parameters
+- Gradient clipping threshold
+- Checkpoint and logging intervals
+
+PACKAGES USED:
+- dataclasses: Clean config structure
+- json: Save/load configs
+
+FILES FROM THIS PROJECT:
+- None (base config)
+
+DESIGN DECISIONS:
+- Small batch size (16-32) for 8GB VRAM
+- Learning rate ~1e-4 to 3e-4 (typical for small transformers)
+- Weight decay 0.01 (L2 regularization)
+- Gradient clipping at 1.0 (prevents exploding gradients)
+- Warmup steps to stabilize early training
+
+COMMON FAILURE MODES:
+- LR too high → loss explodes, NaN
+- LR too low → very slow convergence
+- No warmup → unstable early training
+- Batch size too large → OOM
+- No gradient clipping → gradient explosion
+"""
+
+from dataclasses import dataclass
+from pathlib import Path
+import json
+
+
+@dataclass
+class TrainingConfig:
+    """Training hyperparameters."""
+
+    # Data
+    train_data_path: str = "./data/tokenized/train.npy"
+    val_data_path: str = "./data/tokenized/val.npy"
+
+    # Batch size (adjust based on VRAM)
+    batch_size: int = 16  # GTX 1080: 16-32 works well
+
+    # Training duration
+    num_epochs: int = 10  # Total training epochs
+    max_steps: int = -1  # -1 means train for num_epochs
+
+    # Optimization
+    learning_rate: float = 3e-4  # Peak learning rate
+    weight_decay: float = 0.01  # L2 regularization
+    betas: tuple = (0.9, 0.999)  # Adam beta1, beta2
+    eps: float = 1e-8  # Adam epsilon
+
+    # Learning rate schedule
+    warmup_steps: int = 100  # Linear warmup steps
+    lr_decay: str = "cosine"  # 'cosine' or 'linear' or 'constant'
+    min_lr: float = 1e-5  # Minimum LR for decay
+
+    # Gradient clipping (CRITICAL for stability)
+    grad_clip: float = 1.0  # Clip gradient norm to this value
+
+    # Evaluation
+    eval_interval: int = 500  # Evaluate every N steps
+    eval_iters: int = 100  # Number of eval batches
+
+    # Checkpointing
+    checkpoint_dir: str = "./model/checkpoints"
+    save_interval: int = 1000  # Save checkpoint every N steps
+    keep_last_n: int = 3  # Keep only last N checkpoints
+
+    # Logging
+    log_interval: int = 10  # Log every N steps
+    log_dir: str = "./logs"
+
+    # Device
+    device: str = "cuda"  # 'cuda' or 'cpu'
+
+    # Mixed precision (optional, for faster training)
+    use_amp: bool = False  # Automatic Mixed Precision
+
+    # Reproducibility
+    seed: int = 42
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary."""
+        return {
+            "train_data_path": self.train_data_path,
+            "val_data_path": self.val_data_path,
+            "batch_size": self.batch_size,
+            "num_epochs": self.num_epochs,
+            "max_steps": self.max_steps,
+            "learning_rate": self.learning_rate,
+            "weight_decay": self.weight_decay,
+            "betas": self.betas,
+            "eps": self.eps,
+            "warmup_steps": self.warmup_steps,
+            "lr_decay": self.lr_decay,
+            "min_lr": self.min_lr,
+            "grad_clip": self.grad_clip,
+            "eval_interval": self.eval_interval,
+            "eval_iters": self.eval_iters,
+            "checkpoint_dir": self.checkpoint_dir,
+            "save_interval": self.save_interval,
+            "keep_last_n": self.keep_last_n,
+            "log_interval": self.log_interval,
+            "log_dir": self.log_dir,
+            "device": self.device,
+            "use_amp": self.use_amp,
+            "seed": self.seed,
+        }
+
+    def save(self, path: Path):
+        """Save config to JSON."""
+        with open(path, "w") as f:
+            json.dump(self.to_dict(), f, indent=2)
+
+    @classmethod
+    def from_dict(cls, config_dict: dict):
+        """Load from dictionary."""
+        return cls(**config_dict)
+
+    @classmethod
+    def from_file(cls, path: Path):
+        """Load from JSON file."""
+        with open(path, "r") as f:
+            config_dict = json.load(f)
+        return cls.from_dict(config_dict)
+
+
+def get_default_config() -> TrainingConfig:
+    """Default training config for GTX 1080."""
+    return TrainingConfig()
+
+
+def get_quick_test_config() -> TrainingConfig:
+    """Quick config for testing (small batch, few steps)."""
+    return TrainingConfig(
+        batch_size=4,
+        num_epochs=1,
+        max_steps=100,
+        eval_interval=50,
+        save_interval=50,
+        log_interval=5,
+    )

gptmed/data/__init__.py

@@ -0,0 +1,5 @@
+"""
+Data processing module
+"""
+
+__all__ = ["parsers"]

gptmed/data/parsers/__init__.py

@@ -0,0 +1,10 @@
+"""
+Data parsers for MedQuAD dataset.
+
+This module contains parsers to extract and process medical Q&A pairs
+from various XML sources.
+"""
+
+from .medquad_parser import MedQuADParser
+
+__all__ = ["MedQuADParser"]

gptmed/data/parsers/medquad_parser.py

@@ -0,0 +1,257 @@
+"""
+MedQuAD XML Parser
+
+Parses MedQuAD XML files and extracts question-answer pairs.
+Follows Single Responsibility Principle - handles only XML parsing logic.
+
+Design decisions:
+- Uses lxml for robust XML parsing (handles malformed XML better than xml.etree)
+- Filters empty answers (copyright-removed collections)
+- Preserves question types and focus metadata for future analysis
+- Returns structured data (not formatting) - separation of concerns
+"""
+
+import xml.etree.ElementTree as ET
+from pathlib import Path
+from typing import Dict, List, Optional
+from dataclasses import dataclass
+
+
+@dataclass
+class QAPair:
+    """
+    Structured representation of a question-answer pair.
+
+    Attributes:
+        question: The medical question text
+        answer: The answer text (may be empty for copyright-removed content)
+        qid: Unique question ID
+        qtype: Question type (e.g., 'treatment', 'symptoms', 'causes')
+        focus: Medical entity the question focuses on (disease, drug, etc.)
+        focus_category: Category of focus (Disease, Drug, Other)
+        source: Source collection name
+    """
+
+    question: str
+    answer: str
+    qid: str
+    qtype: str
+    focus: str
+    focus_category: Optional[str]
+    source: str
+
+    def has_answer(self) -> bool:
+        """Check if this pair has a non-empty answer."""
+        return bool(self.answer and self.answer.strip())
+
+    def to_dict(self) -> Dict:
+        """Convert to dictionary for serialization."""
+        return {
+            "question": self.question,
+            "answer": self.answer,
+            "qid": self.qid,
+            "qtype": self.qtype,
+            "focus": self.focus,
+            "focus_category": self.focus_category,
+            "source": self.source,
+        }
+
+
+class MedQuADParser:
+    """
+    Parser for MedQuAD XML files.
+
+    This class is responsible only for parsing XML structure.
+    Formatting to causal text is handled separately (SRP).
+    """
+
+    # Collections known to have removed answers (MedlinePlus copyright)
+    EMPTY_COLLECTIONS = {"10_MPlus_ADAM_QA", "11_MPlusDrugs_QA", "12_MPlusHerbsSupplements_QA"}
+
+    def __init__(self, dataset_root: Path, skip_empty: bool = True):
+        """
+        Initialize parser.
+
+        Args:
+            dataset_root: Path to MedQuAD root directory
+            skip_empty: Whether to skip collections with removed answers
+        """
+        self.dataset_root = Path(dataset_root)
+        self.skip_empty = skip_empty
+
+        if not self.dataset_root.exists():
+            raise ValueError(f"Dataset root does not exist: {dataset_root}")
+
+    def parse_xml_file(self, xml_path: Path) -> List[QAPair]:
+        """
+        Parse a single XML file and extract Q&A pairs.
+
+        Args:
+            xml_path: Path to XML file
+
+        Returns:
+            List of QAPair objects
+
+        Raises:
+            ET.ParseError: If XML is malformed
+        """
+        try:
+            tree = ET.parse(xml_path)
+            root = tree.getroot()
+
+            # Extract document-level metadata
+            source = root.get("source", "Unknown")
+            focus = root.findtext("Focus", default="Unknown")
+
+            # Get focus category if available
+            focus_category = None
+            focus_annotations = root.find("FocusAnnotations")
+            if focus_annotations is not None:
+                category = focus_annotations.find("Category")
+                if category is not None:
+                    focus_category = category.text
+
+            # Extract Q&A pairs
+            qa_pairs = []
+            qa_pairs_elem = root.find("QAPairs")
+
+            if qa_pairs_elem is None:
+                return qa_pairs
+
+            for qa_elem in qa_pairs_elem.findall("QAPair"):
+                question_elem = qa_elem.find("Question")
+                answer_elem = qa_elem.find("Answer")
+
+                if question_elem is None:
+                    continue
+
+                question = question_elem.text or ""
+                answer = answer_elem.text if answer_elem is not None else ""
+                qid = question_elem.get("qid", "")
+                qtype = question_elem.get("qtype", "unknown")
+
+                qa_pair = QAPair(
+                    question=question.strip(),
+                    answer=answer.strip() if answer else "",
+                    qid=qid,
+                    qtype=qtype,
+                    focus=focus,
+                    focus_category=focus_category,
+                    source=source,
+                )
+
+                qa_pairs.append(qa_pair)
+
+            return qa_pairs
+
+        except ET.ParseError as e:
+            raise ET.ParseError(f"Failed to parse {xml_path}: {e}")
+
+    def get_collection_paths(self) -> List[Path]:
+        """
+        Get all collection directories to process.
+
+        Returns:
+            List of collection directory paths
+        """
+        collections = []
+
+        for item in self.dataset_root.iterdir():
+            if not item.is_dir():
+                continue
+
+            # Skip hidden directories and git
+            if item.name.startswith("."):
+                continue
+
+            # Skip empty collections if requested
+            if self.skip_empty and item.name in self.EMPTY_COLLECTIONS:
+                continue
+
+            collections.append(item)
+
+        return sorted(collections)
+
+    def parse_collection(self, collection_path: Path) -> List[QAPair]:
+        """
+        Parse all XML files in a collection directory.
+
+        Args:
+            collection_path: Path to collection directory
+
+        Returns:
+            List of all QAPair objects from this collection
+        """
+        qa_pairs = []
+        xml_files = list(collection_path.glob("*.xml"))
+
+        for xml_file in xml_files:
+            try:
+                pairs = self.parse_xml_file(xml_file)
+                qa_pairs.extend(pairs)
+            except ET.ParseError as e:
+                print(f"Warning: Skipping malformed file {xml_file}: {e}")
+                continue
+
+        return qa_pairs
+
+    def parse_all(self, filter_empty_answers: bool = True) -> List[QAPair]:
+        """
+        Parse all collections in the dataset.
+
+        Args:
+            filter_empty_answers: Whether to filter out pairs with empty answers
+
+        Returns:
+            List of all QAPair objects from all collections
+        """
+        all_pairs = []
+        collections = self.get_collection_paths()
+
+        print(f"Found {len(collections)} collections to process")
+
+        for collection in collections:
+            print(f"Processing {collection.name}...")
+            pairs = self.parse_collection(collection)
+
+            if filter_empty_answers:
+                pairs = [p for p in pairs if p.has_answer()]
+
+            all_pairs.extend(pairs)
+            print(f"  Extracted {len(pairs)} Q&A pairs")
+
+        print(f"\nTotal Q&A pairs: {len(all_pairs)}")
+        return all_pairs
+
+    def get_statistics(self, qa_pairs: List[QAPair]) -> Dict:
+        """
+        Get statistics about parsed Q&A pairs.
+
+        Args:
+            qa_pairs: List of QAPair objects
+
+        Returns:
+            Dictionary with statistics
+        """
+        stats = {
+            "total_pairs": len(qa_pairs),
+            "pairs_with_answers": sum(1 for p in qa_pairs if p.has_answer()),
+            "pairs_without_answers": sum(1 for p in qa_pairs if not p.has_answer()),
+            "sources": {},
+            "qtypes": {},
+            "focus_categories": {},
+        }
+
+        for pair in qa_pairs:
+            # Count by source
+            stats["sources"][pair.source] = stats["sources"].get(pair.source, 0) + 1
+
+            # Count by question type
+            stats["qtypes"][pair.qtype] = stats["qtypes"].get(pair.qtype, 0) + 1
+
+            # Count by focus category
+            if pair.focus_category:
+                cat = pair.focus_category
+                stats["focus_categories"][cat] = stats["focus_categories"].get(cat, 0) + 1
+
+        return stats

gptmed/data/parsers/text_formatter.py

@@ -0,0 +1,148 @@
+"""
+Text Formatter for Causal Language Modeling
+
+Converts structured Q&A pairs into causal text format suitable for
+next-token prediction training.
+
+Design decisions explained:
+- Simple format preserves question-answer structure
+- Special tokens ([Q], [A]) help model learn task boundaries
+- Newlines create clear separation for tokenizer
+- No complex templating - reduces failure modes
+"""
+
+from typing import List
+from dataclasses import dataclass
+
+
+@dataclass
+class FormatConfig:
+    """Configuration for text formatting."""
+
+    use_special_tokens: bool = True  # Use [Q] and [A] markers
+    add_separator: bool = True  # Add newline between Q and A
+    add_end_token: bool = True  # Add end-of-text marker
+    question_prefix: str = "Q: "
+    answer_prefix: str = "A: "
+    separator: str = "\n"
+    end_token: str = "\n\n"  # Double newline = document boundary
+
+
+class CausalTextFormatter:
+    """
+    Formats Q&A pairs for causal language modeling.
+
+    Follows Open-Closed Principle: easy to extend with new formats
+    without modifying existing code.
+    """
+
+    def __init__(self, config: FormatConfig = None):
+        """
+        Initialize formatter.
+
+        Args:
+            config: Formatting configuration
+        """
+        self.config = config or FormatConfig()
+
+    def format_single_pair(self, question: str, answer: str) -> str:
+        """
+        Format a single Q&A pair.
+
+        Args:
+            question: Question text
+            answer: Answer text
+
+        Returns:
+            Formatted text string
+
+        Example:
+            >>> formatter = CausalTextFormatter()
+            >>> formatter.format_single_pair("What is cancer?", "Cancer is...")
+            "Q: What is cancer?\\nA: Cancer is...\\n\\n"
+        """
+        parts = []
+
+        # Add question
+        if self.config.use_special_tokens:
+            parts.append(f"{self.config.question_prefix}{question}")
+        else:
+            parts.append(question)
+
+        # Add separator
+        if self.config.add_separator:
+            parts.append(self.config.separator)
+
+        # Add answer
+        if self.config.use_special_tokens:
+            parts.append(f"{self.config.answer_prefix}{answer}")
+        else:
+            parts.append(answer)
+
+        # Add end token
+        if self.config.add_end_token:
+            parts.append(self.config.end_token)
+
+        return "".join(parts)
+
+    def format_batch(self, qa_pairs: List[tuple]) -> str:
+        """
+        Format multiple Q&A pairs into a single text corpus.
+
+        Args:
+            qa_pairs: List of (question, answer) tuples
+
+        Returns:
+            Single formatted text string
+        """
+        formatted_pairs = [self.format_single_pair(q, a) for q, a in qa_pairs]
+        return "".join(formatted_pairs)
+
+    def format_from_structured(self, qa_objects: List) -> str:
+        """
+        Format from structured QAPair objects.
+
+        Args:
+            qa_objects: List of QAPair objects (from parser)
+
+        Returns:
+            Formatted text corpus
+        """
+        pairs = [(obj.question, obj.answer) for obj in qa_objects]
+        return self.format_batch(pairs)
+
+
+class MinimalFormatter(CausalTextFormatter):
+    """
+    Minimal format without special tokens.
+    Useful for baseline comparisons.
+    """
+
+    def __init__(self):
+        config = FormatConfig(
+            use_special_tokens=False,
+            add_separator=True,
+            add_end_token=True,
+            separator="\n",
+            end_token="\n\n",
+        )
+        super().__init__(config)
+
+
+class StructuredFormatter(CausalTextFormatter):
+    """
+    More structured format with explicit markers.
+    Better for instruction-tuning style training.
+    """
+
+    def __init__(self):
+        config = FormatConfig(
+            use_special_tokens=True,
+            add_separator=True,
+            add_end_token=True,
+            question_prefix="### Question: ",
+            answer_prefix="### Answer: ",
+            separator="\n\n",
+            end_token="\n\n---\n\n",
+        )
+        super().__init__(config)

gptmed/inference/__init__.py

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