feat: support per-sample loss weighting (sample_weights)

docs/source/architecture/overview.md

@@ -569,6 +569,55 @@ class CustomClassifier(nn.Module):
         return self.head(custom_features)
 ```
 
+## Weighting Individual Samples in the Loss (`sample_weights`)
+
+**Purpose:** Make some training (or validation) examples count more than others
+when the loss is computed — useful for class rebalancing, samples with uncertain
+labels, or business-driven importance weighting.
+
+`torchTextClassifiers.train()` accepts an optional `sample_weights` array (one
+weight per training example) and a separate `val_sample_weights` array (one
+weight per validation example). Both default to `1` for every sample, which is
+equivalent to not weighting at all.
+
+```python
+import numpy as np
+
+# Give recent/high-confidence samples more weight, e.g. 2x the default
+sample_weights = np.ones(len(X_train))
+sample_weights[recent_samples_mask] = 2.0
+
+classifier.train(
+    X_train, y_train,
+    training_config=training_config,
+    X_val=X_val, y_val=y_val,
+    sample_weights=sample_weights,          # shape (len(X_train),)
+    val_sample_weights=None,                # validation stays unweighted here
+)
+```
+
+**How it flows through the pipeline:**
+1. `TextClassificationDataset` stores the weights and returns one per sample
+   alongside the tokenized text, categorical variables, and label.
+2. The collate function stacks them into a `sample_weights` tensor on the batch.
+3. `TextClassificationModule` reads `batch["sample_weights"]` and computes a
+   **weighted average loss** instead of a plain mean: standard `torch.nn.*Loss`
+   objects (e.g. `CrossEntropyLoss`, `BCEWithLogitsLoss`) are automatically
+   switched to `reduction="none"` internally so each sample's loss can be
+   scaled by its weight before averaging.
+
+This composes naturally with a loss's own per-*class* `weight=` argument (e.g.
+`torch.nn.CrossEntropyLoss(weight=class_weights)` for class imbalance) — the
+per-sample loss already reflects the class weight, and `sample_weights` is
+applied on top of it.
+
+**Custom losses** (used via `from_model`, see below) support this the same
+way standard PyTorch losses do: `forward` should return an **unreduced,
+per-sample** tensor of shape `(batch,)` (i.e. behave like
+`reduction="none"`), and `TextClassificationModule` takes care of applying
+`sample_weights` and reducing the batch itself — `MultiLevelCrossEntropyLoss`
+in `torchTextClassifiers.contrib` is a working example.
+
 ## Using the High-Level API
 
 For most users, the `torchTextClassifiers` wrapper handles all the complexity:

docs/source/getting_started/quickstart.md

@@ -237,6 +237,33 @@ classifier.train(
 )
 ```
 
+## Weighting Samples in the Loss
+
+If some training examples should count more than others (e.g. to rebalance
+classes or reflect label confidence), pass a `sample_weights` array to
+`train()` — one weight per training sample, defaulting to `1` (no weighting)
+when omitted. A separate `val_sample_weights` array can be passed for the
+validation set:
+
+```python
+import numpy as np
+
+sample_weights = np.ones(len(X_train))
+sample_weights[y_train == 0] = 2.0  # weight the minority class more heavily
+
+classifier.train(
+    X_train=X_train,
+    y_train=y_train,
+    X_val=X_val,
+    y_val=y_val,
+    sample_weights=sample_weights,
+    training_config=training_config,
+)
+```
+
+See {doc}`../architecture/overview` for details on how the weights flow through
+the dataset, batching, and loss computation.
+
 ## What's Next?
 
 Now that you've built your first classifier, you can:

tests/test_sample_weights.py

@@ -0,0 +1,247 @@
+import numpy as np
+import pytest
+import torch
+
+from torchTextClassifiers import ModelConfig, TrainingConfig, torchTextClassifiers
+from torchTextClassifiers.contrib import MultiLevelCrossEntropyLoss
+from torchTextClassifiers.dataset import TextClassificationDataset
+from torchTextClassifiers.model import TextClassificationModule
+from torchTextClassifiers.tokenizers import NGramTokenizer
+
+
+def _trained_ngram_tokenizer(texts):
+    tokenizer = NGramTokenizer(
+        min_count=1, min_n=2, max_n=4, num_tokens=50, len_word_ngrams=2, output_dim=20
+    )
+    tokenizer.train(list(texts))
+    return tokenizer
+
+
+class DummyClassificationModel(torch.nn.Module):
+    """Bypasses tokenization/embedding: forwards pre-computed logits straight through."""
+
+    def __init__(self, num_classes):
+        super().__init__()
+        self.num_classes = num_classes
+        self.categorical_variable_net = None
+
+    def forward(self, input_ids, attention_mask=None, categorical_vars=None, **kwargs):
+        return input_ids
+
+
+class TestDatasetSampleWeights:
+    def test_default_sample_weights_are_ones(self, sample_text_data, sample_labels):
+        tokenizer = _trained_ngram_tokenizer(sample_text_data)
+        dataset = TextClassificationDataset(
+            texts=sample_text_data.tolist(),
+            categorical_variables=None,
+            tokenizer=tokenizer,
+            labels=sample_labels.tolist(),
+        )
+        dataloader = dataset.create_dataloader(
+            batch_size=len(sample_text_data), shuffle=False, num_workers=0
+        )
+        batch = next(iter(dataloader))
+
+        assert torch.allclose(batch["sample_weights"], torch.ones(len(sample_text_data)))
+
+    def test_custom_sample_weights_flow_through_batch(self, sample_text_data, sample_labels):
+        weights = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6], dtype=np.float32)
+        tokenizer = _trained_ngram_tokenizer(sample_text_data)
+        dataset = TextClassificationDataset(
+            texts=sample_text_data.tolist(),
+            categorical_variables=None,
+            tokenizer=tokenizer,
+            labels=sample_labels.tolist(),
+            sample_weights=weights,
+        )
+        dataloader = dataset.create_dataloader(
+            batch_size=len(sample_text_data), shuffle=False, num_workers=0
+        )
+        batch = next(iter(dataloader))
+
+        assert torch.allclose(batch["sample_weights"], torch.tensor(weights))
+
+
+class TestLightningModuleSampleWeights:
+    def _build_module(self, loss, num_classes=3):
+        return TextClassificationModule(
+            model=DummyClassificationModel(num_classes=num_classes),
+            loss=loss,
+            optimizer=torch.optim.Adam,
+            optimizer_params={"lr": 1e-3},
+            scheduler=None,
+            scheduler_params=None,
+        )
+
+    @staticmethod
+    def _make_batch(logits, targets, sample_weights=None):
+        batch = {
+            "input_ids": logits,
+            "attention_mask": None,
+            "categorical_vars": None,
+            "labels": targets,
+        }
+        if sample_weights is not None:
+            batch["sample_weights"] = sample_weights
+        return batch
+
+    def test_default_loss_reduction_switched_to_none(self):
+        module = self._build_module(torch.nn.CrossEntropyLoss())
+        assert module.loss.reduction == "none"
+
+    def test_uniform_weights_match_unweighted_loss(self):
+        torch.manual_seed(0)
+        logits = torch.randn(5, 3)
+        targets = torch.tensor([0, 1, 2, 1, 0])
+
+        module = self._build_module(torch.nn.CrossEntropyLoss())
+        batch = self._make_batch(logits, targets, torch.ones(5))
+        loss, _ = module.step(batch)
+
+        expected = torch.nn.functional.cross_entropy(logits, targets)
+        assert torch.allclose(loss, expected, atol=1e-6)
+
+    def test_missing_sample_weights_defaults_to_ones(self):
+        torch.manual_seed(0)
+        logits = torch.randn(5, 3)
+        targets = torch.tensor([0, 1, 2, 1, 0])
+
+        module = self._build_module(torch.nn.CrossEntropyLoss())
+        batch = self._make_batch(logits, targets)
+        loss, _ = module.step(batch)
+
+        expected = torch.nn.functional.cross_entropy(logits, targets)
+        assert torch.allclose(loss, expected, atol=1e-6)
+
+    def test_weighted_loss_matches_manual_computation(self):
+        torch.manual_seed(0)
+        logits = torch.randn(4, 3)
+        targets = torch.tensor([0, 1, 2, 1])
+        weights = torch.tensor([1.0, 0.0, 2.0, 1.0])
+
+        module = self._build_module(torch.nn.CrossEntropyLoss())
+        batch = self._make_batch(logits, targets, weights)
+        loss, _ = module.step(batch)
+
+        per_sample = torch.nn.functional.cross_entropy(logits, targets, reduction="none")
+        expected = (per_sample * weights).sum() / weights.sum()
+        assert torch.allclose(loss, expected, atol=1e-6)
+
+    def test_zero_weight_equivalent_to_excluding_sample(self):
+        torch.manual_seed(0)
+        logits = torch.randn(3, 3)
+        targets = torch.tensor([0, 1, 2])
+        weights = torch.tensor([1.0, 0.0, 1.0])
+
+        module = self._build_module(torch.nn.CrossEntropyLoss())
+        batch_with_zero = self._make_batch(logits, targets, weights)
+        loss_with_zero, _ = module.step(batch_with_zero)
+
+        kept = [0, 2]
+        batch_excluded = self._make_batch(logits[kept], targets[kept], torch.ones(2))
+        loss_excluded, _ = module.step(batch_excluded)
+
+        assert torch.allclose(loss_with_zero, loss_excluded, atol=1e-6)
+
+
+class TestMultiLevelCrossEntropyLoss:
+    def test_returns_per_sample_tensor(self):
+        torch.manual_seed(0)
+        outputs = [torch.randn(4, 3), torch.randn(4, 2)]
+        labels = torch.stack([torch.tensor([0, 1, 2, 1]), torch.tensor([0, 1, 0, 1])], dim=1)
+
+        loss_fn = MultiLevelCrossEntropyLoss()
+        assert loss_fn.reduction == "none"
+
+        per_sample = loss_fn(outputs, labels)
+        assert per_sample.shape == (4,)
+
+        per_level_losses = [
+            torch.nn.functional.cross_entropy(out, labels[:, i], reduction="none")
+            for i, out in enumerate(outputs)
+        ]
+        expected = sum(per_level_losses) / len(outputs)
+        assert torch.allclose(per_sample, expected, atol=1e-6)
+
+    def test_step_applies_sample_weights_across_levels(self):
+        torch.manual_seed(0)
+        outputs = [torch.randn(4, 3), torch.randn(4, 2)]
+        labels = torch.stack([torch.tensor([0, 1, 2, 1]), torch.tensor([0, 1, 0, 1])], dim=1)
+        weights = torch.tensor([1.0, 0.0, 2.0, 1.0])
+
+        model = DummyClassificationModel(num_classes=[3, 2])
+        module = TextClassificationModule(
+            model=model,
+            loss=MultiLevelCrossEntropyLoss(),
+            optimizer=torch.optim.Adam,
+            optimizer_params={"lr": 1e-3},
+            scheduler=None,
+            scheduler_params=None,
+        )
+        batch = {
+            "input_ids": outputs,
+            "attention_mask": None,
+            "categorical_vars": None,
+            "labels": labels,
+            "sample_weights": weights,
+        }
+        loss, _ = module.step(batch)
+
+        per_level_losses = [
+            torch.nn.functional.cross_entropy(out, labels[:, i], reduction="none")
+            for i, out in enumerate(outputs)
+        ]
+        per_sample = sum(per_level_losses) / len(outputs)
+        expected = (per_sample * weights).sum() / weights.sum()
+
+        assert torch.allclose(loss, expected, atol=1e-6)
+
+
+class TestWrapperSampleWeightsValidation:
+    def test_check_sample_weights_none_passthrough(self):
+        assert torchTextClassifiers._check_sample_weights(None, 5) is None
+
+    def test_check_sample_weights_valid(self):
+        weights = [0.1, 0.2, 0.3]
+        checked = torchTextClassifiers._check_sample_weights(weights, 3)
+        assert isinstance(checked, np.ndarray)
+        assert checked.shape == (3,)
+
+    def test_check_sample_weights_wrong_length_raises(self):
+        with pytest.raises(ValueError):
+            torchTextClassifiers._check_sample_weights([0.1, 0.2], 3)
+
+    def test_check_sample_weights_negative_raises(self):
+        with pytest.raises(AssertionError):
+            torchTextClassifiers._check_sample_weights([-0.1, 0.2, 0.3], 3)
+
+
+class TestTrainWithSampleWeights:
+    def test_train_runs_with_sample_weights(self, sample_text_data, sample_labels):
+        tokenizer = _trained_ngram_tokenizer(sample_text_data)
+        model_config = ModelConfig(embedding_dim=8, num_classes=2)
+        ttc = torchTextClassifiers(tokenizer=tokenizer, model_config=model_config)
+
+        training_config = TrainingConfig(
+            num_epochs=1,
+            batch_size=4,
+            lr=1e-3,
+            num_workers=0,
+            raw_labels=False,
+        )
+
+        sample_weights = np.linspace(0.5, 1.5, num=len(sample_text_data))
+        val_sample_weights = np.linspace(0.5, 1.5, num=len(sample_text_data))
+
+        ttc.train(
+            X_train=sample_text_data,
+            y_train=sample_labels,
+            X_val=sample_text_data,
+            y_val=sample_labels,
+            sample_weights=sample_weights,
+            val_sample_weights=val_sample_weights,
+            training_config=training_config,
+        )
+
+        assert ttc.save_path is not None

torchTextClassifiers/contrib/multilevel.py

@@ -167,10 +167,14 @@ class MultiLevelCrossEntropyLoss(nn.Module):
     def __init__(self, num_classes: Optional[list[int]] = None):
         super().__init__()
         self.num_classes = num_classes
-        self.loss_fn = nn.CrossEntropyLoss()
+        self.reduction = "none"
+        self.loss_fn = nn.CrossEntropyLoss(reduction="none")
 
     def forward(self, outputs: list[torch.Tensor], labels: torch.Tensor) -> torch.Tensor:
-        """Compute the weighted average loss.
+        """Compute the per-sample weighted-average loss across levels.
+
+        Behaves like a ``reduction="none"`` loss: ``TextClassificationModule``
+        applies ``sample_weights`` and reduces the batch itself.
 
         Args:
             outputs: List of logit tensors ``(batch, num_classes_i)`` returned
@@ -179,9 +183,9 @@ def forward(self, outputs: list[torch.Tensor], labels: torch.Tensor) -> torch.Te
                 Column ``i`` contains the ground-truth label for level ``i``.
 
         Returns:
-            Scalar loss tensor.
+            Per-sample loss tensor of shape ``(batch,)``.
         """
-        total_loss = torch.tensor(0.0, device=outputs[0].device)
+        total_loss = torch.zeros(outputs[0].shape[0], device=outputs[0].device)
         for idx, output in enumerate(outputs):
             label = labels[:, idx]
             weight = self.num_classes[idx] if self.num_classes is not None else 1