classone 0.1.0__py3-none-any.whl

classone/__init__.py

@@ -0,0 +1,48 @@
+"""ClassOne: Open-source System 1 Decision Model Architecture built on Gemma 4 E2B."""
+
+from classone.modeling import ClassOneConfig, ClassOneModel
+from classone.schemas import (
+    ChoiceQuestion,
+    ChoiceResult,
+    DecisionRequest,
+    DecisionResponse,
+    DecisionUsage,
+    NoulQuestion,
+    NoulResult,
+    ScoreQuestion,
+    ScoreResult,
+)
+from classone.sdk import (
+    AsyncClassOneClient,
+    Choice,
+    ChoiceAnswer,
+    ClassOneClient,
+    ClassOneResponse,
+    Noul,
+    NoulAnswer,
+    Score,
+    ScoreAnswer,
+)
+
+__all__ = [
+    "Choice",
+    "ChoiceAnswer",
+    "ChoiceQuestion",
+    "ChoiceResult",
+    "ClassOneClient",
+    "AsyncClassOneClient",
+    "ClassOneConfig",
+    "ClassOneModel",
+    "ClassOneResponse",
+    "DecisionRequest",
+    "DecisionResponse",
+    "DecisionUsage",
+    "Noul",
+    "NoulAnswer",
+    "NoulQuestion",
+    "NoulResult",
+    "Score",
+    "ScoreAnswer",
+    "ScoreQuestion",
+    "ScoreResult",
+]

classone/data/__init__.py

@@ -0,0 +1 @@
+"""Data handling and formatting pipelines for ClassOne."""

classone/data/dataset.py

@@ -0,0 +1,175 @@
+"""Dataset pipeline and task formatters for ClassOne System One training.
+
+Provides:
+- Converters for Intent/Routing datasets into Choice questions
+- Converters for Moderation/Verification datasets into Noul questions
+- Converters for Rating/Severity datasets into Score questions
+- Multi-task sample combiners (merging multiple questions per state)
+- PyTorch Dataset and DataLoader integration
+"""
+
+from __future__ import annotations
+
+import json
+import random
+from typing import Any
+
+from torch.utils.data import DataLoader, Dataset
+
+from classone.schemas import ChoiceQuestion, NoulQuestion, Question, ScoreQuestion
+from classone.train.trainer import TrainingItem
+
+
+def create_noul_item(
+    state: str | dict[str, Any] | list[Any],
+    instruction: str,
+    target: bool | float,
+    question_id: str = "q_noul",
+) -> TrainingItem:
+    """Creates a TrainingItem with a single Noul boolean verification question."""
+    target_val = 1.0 if target is True or target == 1 else (0.0 if target is False or target == 0 else float(target))
+    return TrainingItem(
+        state=state,
+        questions={question_id: NoulQuestion(instructions=instruction)},
+        targets={question_id: target_val},
+    )
+
+
+def create_choice_item(
+    state: str | dict[str, Any] | list[Any],
+    instruction: str,
+    target_key: str,
+    all_criteria: dict[str, str],
+    max_options: int | None = None,
+    question_id: str = "q_choice",
+) -> TrainingItem:
+    """Creates a TrainingItem with a Choice classification question.
+
+    If max_options is specified and smaller than total categories, subsamples
+    distractor options while always including the target_key.
+    """
+    if target_key not in all_criteria:
+        raise ValueError(f"target_key '{target_key}' not found in all_criteria keys.")
+
+    if max_options and max_options < len(all_criteria):
+        distractor_keys = [k for k in all_criteria if k != target_key]
+        num_distractors = min(len(distractor_keys), max(1, max_options - 1))
+        sampled_keys = random.sample(distractor_keys, num_distractors) + [target_key]
+        random.shuffle(sampled_keys)
+        criteria = {k: all_criteria[k] for k in sampled_keys}
+    else:
+        criteria = dict(all_criteria)
+
+    return TrainingItem(
+        state=state,
+        questions={
+            question_id: ChoiceQuestion(
+                instructions=instruction,
+                criteria=criteria,
+            )
+        },
+        targets={question_id: target_key},
+    )
+
+
+def create_score_item(
+    state: str | dict[str, Any] | list[Any],
+    instruction: str,
+    target_level: float | str,
+    rubric: list[str],
+    question_id: str = "q_score",
+) -> TrainingItem:
+    """Creates a TrainingItem with an ordinal Score rubric question."""
+    return TrainingItem(
+        state=state,
+        questions={
+            question_id: ScoreQuestion(
+                instructions=instruction,
+                criteria=rubric,
+            )
+        },
+        targets={question_id: target_level},
+    )
+
+
+def bundle_multi_task_items(
+    state: str | dict[str, Any] | list[Any],
+    questions: dict[str, Question],
+    targets: dict[str, int | float | str],
+) -> TrainingItem:
+    """Bundles multiple parallel questions over a single shared input state."""
+    return TrainingItem(state=state, questions=questions, targets=targets)
+
+
+class ClassOneDataset(Dataset):
+    """PyTorch Dataset containing ClassOne TrainingItems."""
+
+    def __init__(self, items: list[TrainingItem]):
+        self.items = items
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __getitem__(self, idx: int) -> TrainingItem:
+        return self.items[idx]
+
+    def save_jsonl(self, path: str):
+        """Exports dataset items to JSONL format."""
+        with open(path, "w", encoding="utf-8") as f:
+            for item in self.items:
+                serialized_questions = {}
+                for q_id, q in item.questions.items():
+                    serialized_questions[q_id] = q.model_dump()
+
+                row = {
+                    "state": item.state,
+                    "questions": serialized_questions,
+                    "targets": item.targets,
+                }
+                f.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+    @classmethod
+    def load_jsonl(cls, path: str) -> ClassOneDataset:
+        """Loads dataset from JSONL format."""
+        items: list[TrainingItem] = []
+        with open(path, encoding="utf-8") as f:
+            for line in f:
+                if not line.strip():
+                    continue
+                row = json.loads(line)
+                questions: dict[str, Question] = {}
+                for q_id, q_data in row["questions"].items():
+                    q_type = q_data.get("type")
+                    if q_type == "noul":
+                        questions[q_id] = NoulQuestion(**q_data)
+                    elif q_type == "choice":
+                        questions[q_id] = ChoiceQuestion(**q_data)
+                    elif q_type == "score":
+                        questions[q_id] = ScoreQuestion(**q_data)
+                items.append(
+                    TrainingItem(
+                        state=row["state"],
+                        questions=questions,
+                        targets=row["targets"],
+                    )
+                )
+        return cls(items)
+
+
+def classone_collate_fn(batch: list[TrainingItem]) -> list[TrainingItem]:
+    """Collate function for PyTorch DataLoader preserving List[TrainingItem]."""
+    return batch
+
+
+def create_classone_dataloader(
+    dataset: ClassOneDataset,
+    batch_size: int = 4,
+    shuffle: bool = True,
+) -> DataLoader:
+    """Creates a PyTorch DataLoader yielding batches of TrainingItem."""
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        collate_fn=classone_collate_fn,
+    )

classone/modeling/__init__.py

@@ -0,0 +1,27 @@
+"""Modeling components for ClassOne System 1 decision models."""
+
+from classone.modeling.configuration_classone import ClassOneConfig
+from classone.modeling.loss import (
+    RLCDLoss,
+    brier_score_binary,
+    brier_score_multiclass,
+    expected_calibration_error,
+)
+from classone.modeling.modeling_classone import (
+    ChoiceHead,
+    ClassOneModel,
+    NoulHead,
+    ScoreHead,
+)
+
+__all__ = [
+    "ClassOneModel",
+    "ClassOneConfig",
+    "NoulHead",
+    "ChoiceHead",
+    "ScoreHead",
+    "RLCDLoss",
+    "brier_score_binary",
+    "brier_score_multiclass",
+    "expected_calibration_error",
+]

classone/modeling/configuration_classone.py

@@ -0,0 +1,29 @@
+"""Configuration for ClassOne System 1 decision models."""
+
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+class ClassOneConfig(PretrainedConfig):
+    """Configuration class for ClassOne decision model."""
+
+    model_type = "classone"
+
+    def __init__(
+        self,
+        base_model_name_or_path: str = "google/gemma-2-2b-it",
+        hidden_size: int = 2048,
+        head_hidden_size: int = 512,
+        dropout: float = 0.1,
+        temperature: float = 1.0,
+        use_temperature_scaling: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.base_model_name_or_path = base_model_name_or_path
+        self.hidden_size = hidden_size
+        self.head_hidden_size = head_hidden_size
+        self.dropout = dropout
+        self.temperature = temperature
+        self.use_temperature_scaling = use_temperature_scaling

classone/modeling/loss.py

@@ -0,0 +1,131 @@
+"""Proper scoring rules and calibration loss functions for RLCD.
+
+Implements:
+- Brier Score Loss (normalized mean squared error of probability distributions)
+  Reference: Brier, G. W. (1950). Verification of forecasts expressed in terms of probability.
+  Reference: Gneiting, T., & Raftery, A. E. (2007). Strictly proper scoring rules, prediction, and estimation.
+- Logarithmic Scoring Rule (Cross-Entropy / Negative Log-Likelihood)
+- Expected Calibration Error (ECE)
+  Reference: Guo, C., Pleiss, G., Sun, Y., & Weinberger, K. Q. (2017). On calibration of modern neural networks.
+- Hybrid RLCD Loss combining log-loss and Brier penalty
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+def brier_score_binary(probs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+    """Computes binary Brier score: (p - y)^2 in [0, 1].
+
+    Args:
+        probs: Predicted probability tensor of shape (N,) or (N, 1), in [0, 1].
+        targets: Binary label tensor of shape (N,) or (N, 1), in {0, 1}.
+    """
+    return torch.mean((probs.view(-1).float() - targets.view(-1).float()) ** 2)
+
+
+def brier_score_multiclass(probs: torch.Tensor, targets: torch.Tensor, num_classes: int) -> torch.Tensor:
+    """Computes normalized multiclass Brier score: 0.5 * sum_k (p_k - y_k)^2 in [0, 1].
+
+    Dividing by 2 is the standard normalization in modern probabilistic forecasting
+    and machine learning (e.g. Gneiting & Raftery 2007), ensuring that for K=2 classes
+    it reduces identically to binary Brier score (p - y)^2.
+
+    Args:
+        probs: Probability distribution tensor of shape (N, K).
+        targets: Ground truth class index tensor of shape (N,).
+        num_classes: Number of candidate classes K.
+    """
+    probs_f = probs.float()
+    one_hot = F.one_hot(targets, num_classes=num_classes).float()
+    return 0.5 * torch.mean(torch.sum((probs_f - one_hot) ** 2, dim=-1))
+
+
+def expected_calibration_error(probs: torch.Tensor, targets: torch.Tensor, n_bins: int = 10) -> float:
+    """Calculates the Expected Calibration Error (ECE) across confidence bins.
+
+    Reference: Guo et al. (ICML 2017) "On Calibration of Modern Neural Networks"
+
+    Args:
+        probs: Model confidence/probability for the predicted class, shape (N,).
+        targets: 1 if prediction is correct, 0 otherwise, shape (N,).
+        n_bins: Number of confidence bins.
+    """
+    probs = probs.view(-1).detach().cpu().float()
+    targets = targets.view(-1).detach().cpu().float()
+
+    n_samples = len(probs)
+    if n_samples == 0:
+        return 0.0
+
+    bin_boundaries = torch.linspace(0, 1, n_bins + 1)
+    ece = 0.0
+
+    for i in range(n_bins):
+        bin_lower = bin_boundaries[i]
+        bin_upper = bin_boundaries[i + 1]
+
+        # For the first bin, include 0.0 with >= to prevent exclusion of zero predictions
+        if i == 0:
+            in_bin = (probs >= bin_lower) & (probs <= bin_upper)
+        else:
+            in_bin = (probs > bin_lower) & (probs <= bin_upper)
+
+        count_in_bin = in_bin.sum().item()
+        if count_in_bin > 0:
+            accuracy_in_bin = targets[in_bin].mean().item()
+            avg_confidence_in_bin = probs[in_bin].mean().item()
+            ece += abs(avg_confidence_in_bin - accuracy_in_bin) * (count_in_bin / n_samples)
+
+    return float(ece)
+
+
+class RLCDLoss(nn.Module):
+    """Reinforcement Learning for Calibrated Decisions (RLCD) multi-objective loss.
+
+    Combines Logarithmic Scoring (proper scoring rule rewarding true likelihood)
+    with Brier Score (penalizing probability deviation).
+    """
+
+    def __init__(
+        self,
+        brier_weight: float = 0.5,
+        log_weight: float = 1.0,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+        self.brier_weight = brier_weight
+        self.log_weight = log_weight
+        self.eps = eps
+
+    def forward_binary(self, prob: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        """Loss for Noul boolean decisions."""
+        prob_f = prob.view(-1).float()
+        target_f = target.view(-1).float()
+        prob_clamped = torch.clamp(prob_f, self.eps, 1.0 - self.eps)
+
+        # Binary log loss
+        log_loss = -(target_f * torch.log(prob_clamped) + (1.0 - target_f) * torch.log(1.0 - prob_clamped)).mean()
+
+        # Un-clamped Brier loss for pure proper scoring
+        brier = torch.mean((prob_f - target_f) ** 2)
+
+        return self.log_weight * log_loss + self.brier_weight * brier
+
+    def forward_multiclass(self, probs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        """Loss for Choice / Score decisions."""
+        num_classes = probs.shape[-1]
+        probs_f = probs.float()
+        one_hot = F.one_hot(targets, num_classes=num_classes).float()
+
+        # Log loss computed with stable FP32 clamping
+        probs_clamped = torch.clamp(probs_f, self.eps, 1.0 - self.eps)
+        log_loss = -torch.sum(one_hot * torch.log(probs_clamped), dim=-1).mean()
+
+        # Normalized multiclass Brier score (0.5 * sum squared difference)
+        brier = 0.5 * torch.mean(torch.sum((probs_f - one_hot) ** 2, dim=-1))
+
+        return self.log_weight * log_loss + self.brier_weight * brier