cotlab 0.8.0__py3-none-any.whl

cotlab/core/config.py

@@ -0,0 +1,90 @@
+"""Hydra-compatible configuration dataclasses."""
+
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+from omegaconf import MISSING
+
+
+@dataclass
+class BackendConfig:
+    """Configuration for inference backend."""
+
+    _target_: str = MISSING
+    device: str = "cuda"
+    dtype: str = "bfloat16"
+
+
+@dataclass
+class TransformersBackendConfig(BackendConfig):
+    """Transformers-specific backend config."""
+
+    _target_: str = "cotlab.backends.TransformersBackend"
+    enable_hooks: bool = True
+    trust_remote_code: bool = True
+
+
+@dataclass
+class VLLMBackendConfig(BackendConfig):
+    """vLLM-specific backend config."""
+
+    _target_: str = "cotlab.backends.VLLMBackend"
+    tensor_parallel_size: int = 1
+    max_model_len: int = 4096
+    trust_remote_code: bool = True
+
+
+@dataclass
+class ModelConfig:
+    """Configuration for model loading."""
+
+    name: str = MISSING
+    variant: str = "4b"
+    max_new_tokens: int = 512
+    temperature: float = 0.7
+    top_p: float = 0.9
+    do_sample: bool = True
+
+
+@dataclass
+class PromptConfig:
+    """Configuration for prompt strategy."""
+
+    _target_: str = MISSING
+    name: str = MISSING
+    system_role: Optional[str] = None
+
+
+@dataclass
+class DatasetConfig:
+    """Configuration for dataset loading."""
+
+    _target_: str = MISSING
+    name: str = MISSING
+    path: str = MISSING
+
+
+@dataclass
+class ExperimentConfig:
+    """Configuration for an experiment."""
+
+    _target_: str = MISSING
+    name: str = MISSING
+    description: str = ""
+    num_samples: Optional[int] = None  # None = use all available samples
+    tests: List[str] = field(default_factory=list)
+    metrics: List[str] = field(default_factory=list)
+
+
+@dataclass
+class Config:
+    """Root configuration."""
+
+    backend: BackendConfig = field(default_factory=BackendConfig)
+    model: ModelConfig = field(default_factory=ModelConfig)
+    prompt: PromptConfig = field(default_factory=PromptConfig)
+    dataset: DatasetConfig = field(default_factory=DatasetConfig)
+    experiment: ExperimentConfig = field(default_factory=ExperimentConfig)
+    seed: int = 42
+    verbose: bool = True
+    dry_run: bool = False

cotlab/core/registry.py

@@ -0,0 +1,68 @@
+"""Component registry for dynamic instantiation."""
+
+from typing import Any, Dict, Type
+
+from hydra.utils import instantiate
+from omegaconf import DictConfig
+
+
+class Registry:
+    """Registry for dynamically instantiating components from config."""
+
+    _backends: Dict[str, Type] = {}
+    _prompts: Dict[str, Type] = {}
+    _experiments: Dict[str, Type] = {}
+    _datasets: Dict[str, Type] = {}
+
+    @classmethod
+    def register_backend(cls, name: str):
+        """Decorator to register a backend class."""
+
+        def decorator(klass):
+            cls._backends[name] = klass
+            return klass
+
+        return decorator
+
+    @classmethod
+    def register_prompt(cls, name: str):
+        """Decorator to register a prompt strategy class."""
+
+        def decorator(klass):
+            cls._prompts[name] = klass
+            return klass
+
+        return decorator
+
+    @classmethod
+    def register_experiment(cls, name: str):
+        """Decorator to register an experiment class."""
+
+        def decorator(klass):
+            cls._experiments[name] = klass
+            return klass
+
+        return decorator
+
+    @classmethod
+    def register_dataset(cls, name: str):
+        """Decorator to register a dataset class."""
+
+        def decorator(klass):
+            cls._datasets[name] = klass
+            return klass
+
+        return decorator
+
+
+def create_component(cfg: DictConfig) -> Any:
+    """
+    Create a component from a Hydra config using _target_.
+
+    Args:
+        cfg: DictConfig with _target_ specifying the class
+
+    Returns:
+        Instantiated component
+    """
+    return instantiate(cfg)

cotlab/datasets/__init__.py

@@ -0,0 +1,45 @@
+"""Datasets module."""
+
+from .loaders import (
+    BaseDataset,
+    CardiologyDataset,
+    HistopathologyDataset,
+    MARCDataset,
+    MedBulletsDataset,
+    MedQADataset,
+    MMLUMedicalDataset,
+    NeurologyDataset,
+    OncologyDataset,
+    PatchingPairsDataset,
+    PediatricsDataset,
+    PLABDataset,
+    ProbingDiagnosisDataset,
+    PubHealthBenchDataset,
+    PubMedQADataset,
+    RadiologyDataset,
+    Sample,
+    SyntheticMedicalDataset,
+    TCGADataset,
+)
+
+__all__ = [
+    "Sample",
+    "BaseDataset",
+    "CardiologyDataset",
+    "HistopathologyDataset",
+    "MARCDataset",
+    "MedBulletsDataset",
+    "MedQADataset",
+    "MMLUMedicalDataset",
+    "NeurologyDataset",
+    "OncologyDataset",
+    "RadiologyDataset",
+    "PediatricsDataset",
+    "PLABDataset",
+    "PubMedQADataset",
+    "PubHealthBenchDataset",
+    "SyntheticMedicalDataset",
+    "PatchingPairsDataset",
+    "ProbingDiagnosisDataset",
+    "TCGADataset",
+]