From 5eb584b9f30bcb4e31973181c19cdd13a68a44f0 Mon Sep 17 00:00:00 2001
From: meilame-tayebjee <meilame.tayebjee@ensae.fr>
Date: Thu, 2 Jul 2026 14:17:44 +0000
Subject: [PATCH 1/6] feat(dataset): support per-sample loss weights in
 TextClassificationDataset

Add an optional sample_weights array (defaulting to 1 for every sample)
that flows through __getitem__ and collate_fn into a new "sample_weights"
key in the batch dict, so downstream loss computation can weight samples
individually.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 torchTextClassifiers/dataset/dataset.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/torchTextClassifiers/dataset/dataset.py b/torchTextClassifiers/dataset/dataset.py
index 064d78b..44a1254 100644
--- a/torchTextClassifiers/dataset/dataset.py
+++ b/torchTextClassifiers/dataset/dataset.py
@@ -20,9 +20,16 @@ def __init__(
         tokenizer: BaseTokenizer,
         labels: Union[List[int], List[List[int]], np.array, None] = None,
         ragged_multilabel: bool = False,
+        sample_weights: Union[List[float], np.ndarray, None] = None,
     ):
         self.categorical_variables = categorical_variables
 
+        self.sample_weights = (
+            torch.tensor(sample_weights, dtype=torch.float32)
+            if sample_weights is not None
+            else None
+        )
+
         self.texts = texts
 
         if hasattr(tokenizer, "trained") and not tokenizer.trained:
@@ -59,6 +66,7 @@ def __len__(self):
         return len(self.texts)
 
     def __getitem__(self, idx):
+        weight = self.sample_weights[idx] if self.sample_weights is not None else 1.0
         if self.labels is not None:
             return (
                 str(self.texts[idx]),
@@ -68,6 +76,7 @@ def __getitem__(self, idx):
                     else None
                 ),
                 self.labels[idx],
+                weight,
             )
         else:
             return (
@@ -78,10 +87,11 @@ def __getitem__(self, idx):
                     else None
                 ),
                 None,
+                weight,
             )
 
     def collate_fn(self, batch):
-        text, *categorical_vars, labels = zip(*batch)
+        text, *categorical_vars, labels, weights = zip(*batch)
 
         if self.labels is not None:
             if self.ragged_multilabel:
@@ -120,11 +130,14 @@ def collate_fn(self, batch):
         else:
             categorical_tensors = None
 
+        weights_tensor = torch.tensor([float(w) for w in weights], dtype=torch.float32)
+
         return {
             "input_ids": tokenize_output.input_ids,
             "attention_mask": tokenize_output.attention_mask,
             "categorical_vars": categorical_tensors,
             "labels": labels_tensor,
+            "sample_weights": weights_tensor,
         }
 
     def create_dataloader(

From 471356d9d0b4cf9e7ca951c872b6190c86e324de Mon Sep 17 00:00:00 2001
From: meilame-tayebjee <meilame.tayebjee@ensae.fr>
Date: Thu, 2 Jul 2026 14:17:56 +0000
Subject: [PATCH 2/6] feat(model): weight the loss per sample_weights during
 training

TextClassificationModule now reads sample_weights from the batch and
applies them when computing the loss: losses that accept a
sample_weights kwarg (e.g. MultiLevelCrossEntropyLoss) receive it
directly, otherwise reduction is switched to "none" and the module
computes the weighted mean itself. With all weights equal to 1 this is
mathematically identical to the previous unweighted reduction.

MultiLevelCrossEntropyLoss.forward gains an optional sample_weights
argument, weighting each level's per-sample loss before combining them.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 torchTextClassifiers/contrib/multilevel.py | 18 ++++++++++---
 torchTextClassifiers/model/lightning.py    | 30 +++++++++++++++++++++-
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/torchTextClassifiers/contrib/multilevel.py b/torchTextClassifiers/contrib/multilevel.py
index 7a39755..9db1752 100644
--- a/torchTextClassifiers/contrib/multilevel.py
+++ b/torchTextClassifiers/contrib/multilevel.py
@@ -167,9 +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.loss_fn = nn.CrossEntropyLoss(reduction="none")
 
-    def forward(self, outputs: list[torch.Tensor], labels: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self,
+        outputs: list[torch.Tensor],
+        labels: torch.Tensor,
+        sample_weights: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
         """Compute the weighted average loss.
 
         Args:
@@ -177,15 +182,22 @@ def forward(self, outputs: list[torch.Tensor], labels: torch.Tensor) -> torch.Te
                 by ``MultiLevelTextClassificationModel``.
             labels: Integer label tensor of shape ``(batch, n_levels)``.
                 Column ``i`` contains the ground-truth label for level ``i``.
+            sample_weights: Optional per-sample weight tensor of shape
+                ``(batch,)``. If ``None``, all samples are weighted equally.
 
         Returns:
             Scalar loss tensor.
         """
+        if sample_weights is None:
+            sample_weights = torch.ones(outputs[0].shape[0], device=outputs[0].device)
+
         total_loss = torch.tensor(0.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
-            total_loss = total_loss + self.loss_fn(output.squeeze(), label) * weight
+            per_sample_loss = self.loss_fn(output.squeeze(), label)
+            level_loss = (per_sample_loss * sample_weights).sum() / sample_weights.sum()
+            total_loss = total_loss + level_loss * weight
 
         total_weight = sum(self.num_classes) if self.num_classes is not None else len(outputs)
         return total_loss / total_weight
diff --git a/torchTextClassifiers/model/lightning.py b/torchTextClassifiers/model/lightning.py
index 1b9ed94..5c38e8b 100644
--- a/torchTextClassifiers/model/lightning.py
+++ b/torchTextClassifiers/model/lightning.py
@@ -1,8 +1,13 @@
+import inspect
+import logging
+
 import pytorch_lightning as pl
 import torch
 from torch import nn
 from torchmetrics import Accuracy
 
+logger = logging.getLogger(__name__)
+
 # ============================================================================
 # PyTorch Lightning Module
 # ============================================================================
@@ -40,6 +45,17 @@ def __init__(
         self.model = model
         self.loss = loss
 
+        self._loss_accepts_sample_weights = "sample_weights" in inspect.signature(
+            self.loss.forward
+        ).parameters
+        if not self._loss_accepts_sample_weights and hasattr(self.loss, "reduction"):
+            if self.loss.reduction != "none":
+                logger.info(
+                    f"Setting reduction='none' on {type(self.loss).__name__} so that "
+                    "sample_weights can be applied per-sample before averaging."
+                )
+                self.loss.reduction = "none"
+
         if not hasattr(self.model, "num_classes") or self.model.num_classes is None:
             raise ValueError("Model must have num_classes attribute for accuracy calculation.")
 
@@ -78,7 +94,19 @@ def step(self, batch) -> tuple[torch.Tensor, torch.Tensor | list[torch.Tensor]]:
         outputs = self.forward(batch)
         if isinstance(self.loss, torch.nn.BCEWithLogitsLoss):
             targets = targets.float()
-        loss = self.loss(outputs, targets)
+
+        sample_weights = batch.get("sample_weights")
+        if sample_weights is None:
+            sample_weights = torch.ones(targets.shape[0], device=targets.device)
+        sample_weights = sample_weights.to(targets.device)
+
+        if self._loss_accepts_sample_weights:
+            loss = self.loss(outputs, targets, sample_weights=sample_weights)
+        else:
+            per_sample_loss = self.loss(outputs, targets)
+            per_sample_loss = per_sample_loss.reshape(per_sample_loss.size(0), -1).mean(dim=1)
+            loss = (per_sample_loss * sample_weights).sum() / sample_weights.sum()
+
         if self.multilevel_accuracy:
             accuracy = [
                 fn(out, targets[:, i]) for i, (fn, out) in enumerate(zip(self.accuracy_fn, outputs))

From 06c10905ddefc6aa84958086ea4b72fd925c8822 Mon Sep 17 00:00:00 2001
From: meilame-tayebjee <meilame.tayebjee@ensae.fr>
Date: Thu, 2 Jul 2026 14:18:08 +0000
Subject: [PATCH 3/6] feat(torchTextClassifiers): expose sample_weights and
 val_sample_weights on train()

Add optional sample_weights/val_sample_weights parameters to train(),
validated via a new _check_sample_weights helper (1D array, matching
length, non-negative) and forwarded to the train/val
TextClassificationDataset instances. Also documents the sample_weights
contract for custom losses used via from_model().

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 torchTextClassifiers/torchTextClassifiers.py | 44 ++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/torchTextClassifiers/torchTextClassifiers.py b/torchTextClassifiers/torchTextClassifiers.py
index cd0ce05..3200aef 100644
--- a/torchTextClassifiers/torchTextClassifiers.py
+++ b/torchTextClassifiers/torchTextClassifiers.py
@@ -281,6 +281,18 @@ def forward(
            The wrapper reads ``categorical_variable_net.categorical_vocabulary_sizes``
            to set up the data pipeline.
 
+        **Optional: ``sample_weights`` support in a custom loss** — to support
+        per-sample loss weighting (see ``torchTextClassifiers.train``'s
+        ``sample_weights``/``val_sample_weights`` arguments) with a custom
+        multi-task loss, add an optional ``sample_weights`` keyword argument to
+        its ``forward`` method, e.g. ``forward(self, outputs, labels,
+        sample_weights=None)``, and use it to compute a weighted average
+        instead of a plain mean (see ``MultiLevelCrossEntropyLoss`` below for
+        an example). If your loss exposes a ``reduction`` attribute instead
+        (like standard ``torch.nn.*Loss`` classes) it will automatically be
+        switched to ``"none"`` so that the wrapper can apply the weights and
+        reduce the loss itself.
+
         See ``torchTextClassifiers.contrib`` for ready-made example architectures
         (``MultiLevelTextClassificationModel``, ``MultiLevelCrossEntropyLoss``) that
         follow this interface.
@@ -323,6 +335,8 @@ def train(
         training_config: TrainingConfig,
         X_val: Optional[np.ndarray] = None,
         y_val: Optional[np.ndarray] = None,
+        sample_weights: Optional[np.ndarray] = None,
+        val_sample_weights: Optional[np.ndarray] = None,
         verbose: bool = False,
     ) -> None:
         """Train the classifier using PyTorch Lightning.
@@ -346,6 +360,12 @@ def train(
             X_val: Validation input data
             y_val: Validation labels
             training_config: Configuration parameters for training
+            sample_weights: Optional per-sample weights for the training loss,
+                as a 1D array of length ``len(X_train)``. Defaults to 1 for
+                every sample (i.e. no weighting) when not provided.
+            val_sample_weights: Optional per-sample weights for the
+                validation loss, as a 1D array of length ``len(X_val)``.
+                Defaults to 1 for every sample when not provided.
             verbose: Whether to print training progress information
 
 
@@ -369,6 +389,7 @@ def train(
         X_train, y_train = self._check_XY(
             X_train, y_train, training_config.raw_categorical_inputs, training_config.raw_labels
         )
+        sample_weights = self._check_sample_weights(sample_weights, X_train["text"].shape[0])
 
         if X_val is not None:
             assert y_val is not None, "y_val must be provided if X_val is provided."
@@ -378,6 +399,9 @@ def train(
         X_val_checked: Optional[Dict[str, Any]] = None
         if X_val is not None and y_val is not None:
             X_val_checked, y_val = self._check_XY(X_val, y_val, training_config.raw_categorical_inputs, training_config.raw_labels)
+            val_sample_weights = self._check_sample_weights(
+                val_sample_weights, X_val_checked["text"].shape[0]
+            )
         X_val = X_val_checked
 
         if (
@@ -434,6 +458,7 @@ def train(
             tokenizer=self.tokenizer,
             labels=y_train.tolist(),
             ragged_multilabel=self.ragged_multilabel,
+            sample_weights=sample_weights,
         )
         train_dataloader = train_dataset.create_dataloader(
             batch_size=training_config.batch_size,
@@ -449,6 +474,7 @@ def train(
                 tokenizer=self.tokenizer,
                 labels=y_val,
                 ragged_multilabel=self.ragged_multilabel,
+                sample_weights=val_sample_weights,
             )
             val_dataloader = val_dataset.create_dataloader(
                 batch_size=training_config.batch_size,
@@ -531,6 +557,24 @@ def _check_XY(
 
         return X_checked, Y_checked
 
+    @staticmethod
+    def _check_sample_weights(
+        sample_weights: Optional[np.ndarray], n_samples: int
+    ) -> Optional[np.ndarray]:
+        if sample_weights is None:
+            return None
+
+        sample_weights = np.asarray(sample_weights, dtype=np.float32)
+        assert sample_weights.ndim == 1, "sample_weights must be a 1D array."
+        if sample_weights.shape[0] != n_samples:
+            raise ValueError(
+                f"sample_weights must have length {n_samples} (one weight per sample), "
+                f"got {sample_weights.shape[0]}."
+            )
+        assert (sample_weights >= 0).all(), "sample_weights must be non-negative."
+
+        return sample_weights
+
     @staticmethod
     def _check_text_col(X):
         assert isinstance(

From e120d77002df1e0581d4fe806d38addf51437a33 Mon Sep 17 00:00:00 2001
From: meilame-tayebjee <meilame.tayebjee@ensae.fr>
Date: Thu, 2 Jul 2026 14:18:15 +0000
Subject: [PATCH 4/6] test: cover sample_weights across dataset, lightning
 module, and wrapper

Add unit tests for default/custom sample_weights in
TextClassificationDataset's collate output, weighted-loss computation
in TextClassificationModule.step (including zero-weight equivalent to
sample exclusion), MultiLevelCrossEntropyLoss weighting, wrapper-level
validation, and an end-to-end train() run with sample_weights and
val_sample_weights.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 tests/test_sample_weights.py | 224 +++++++++++++++++++++++++++++++++++
 1 file changed, 224 insertions(+)
 create mode 100644 tests/test_sample_weights.py

diff --git a/tests/test_sample_weights.py b/tests/test_sample_weights.py
new file mode 100644
index 0000000..9de8550
--- /dev/null
+++ b/tests/test_sample_weights.py
@@ -0,0 +1,224 @@
+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 TestMultiLevelCrossEntropyLossSampleWeights:
+    def test_weighted_matches_manual_computation(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])
+
+        loss_fn = MultiLevelCrossEntropyLoss()
+        loss = loss_fn(outputs, labels, sample_weights=weights)
+
+        per_level_losses = []
+        for i, out in enumerate(outputs):
+            per_sample = torch.nn.functional.cross_entropy(out, labels[:, i], reduction="none")
+            per_level_losses.append((per_sample * weights).sum() / weights.sum())
+        expected = sum(per_level_losses) / len(outputs)
+
+        assert torch.allclose(loss, expected, atol=1e-6)
+
+    def test_none_sample_weights_matches_unweighted(self):
+        torch.manual_seed(0)
+        outputs = [torch.randn(4, 3)]
+        labels = torch.tensor([0, 1, 2, 1]).unsqueeze(1)
+
+        loss_fn = MultiLevelCrossEntropyLoss()
+        weighted = loss_fn(outputs, labels)
+        expected = torch.nn.functional.cross_entropy(outputs[0], labels[:, 0])
+
+        assert torch.allclose(weighted, 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

From 9967f6940c6d99b233c80d6995e1ba1b58942e05 Mon Sep 17 00:00:00 2001
From: meilame-tayebjee <meilame.tayebjee@ensae.fr>
Date: Thu, 2 Jul 2026 14:24:13 +0000
Subject: [PATCH 5/6] docs: document sample_weights / val_sample_weights

Add a "Weighting Individual Samples in the Loss" section to the
architecture overview explaining how sample_weights flows through the
dataset, batching, and loss computation, plus a shorter practical
example in the quickstart guide.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 docs/source/architecture/overview.md      | 47 +++++++++++++++++++++++
 docs/source/getting_started/quickstart.md | 27 +++++++++++++
 2 files changed, 74 insertions(+)

diff --git a/docs/source/architecture/overview.md b/docs/source/architecture/overview.md
index e0b5e84..34aad06 100644
--- a/docs/source/architecture/overview.md
+++ b/docs/source/architecture/overview.md
@@ -569,6 +569,53 @@ 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) can opt in by adding an
+optional `sample_weights` keyword argument to `forward`, e.g.
+`forward(self, outputs, labels, sample_weights=None)` — `MultiLevelCrossEntropyLoss`
+in `torchTextClassifiers.contrib` is a working example.
+
 ## Using the High-Level API
 
 For most users, the `torchTextClassifiers` wrapper handles all the complexity:
diff --git a/docs/source/getting_started/quickstart.md b/docs/source/getting_started/quickstart.md
index c6fe663..a80b508 100644
--- a/docs/source/getting_started/quickstart.md
+++ b/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:

From d99de60d7f0c2925f5bb26705b55f72609bf566e Mon Sep 17 00:00:00 2001
From: meilame-tayebjee <meilame.tayebjee@ensae.fr>
Date: Thu, 2 Jul 2026 14:37:32 +0000
Subject: [PATCH 6/6] refactor(model): drop inspect-based loss detection for
 sample_weights

Replace the inspect.signature probing of whether a loss accepts a
sample_weights kwarg with a single, classic convention: any loss used
with TextClassificationModule must expose per-sample (reduction="none")
output. Standard torch.nn.*Loss objects are switched to reduction="none"
automatically; custom losses (e.g. MultiLevelCrossEntropyLoss) just
return an unreduced (batch,) tensor directly. TextClassificationModule
always applies sample_weights and reduces the batch itself, with no
branching. This also keeps torch.nn.CrossEntropyLoss's own per-class
weight= argument untouched and fully compatible.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 docs/source/architecture/overview.md         |  8 +--
 tests/test_sample_weights.py                 | 57 ++++++++++++++------
 torchTextClassifiers/contrib/multilevel.py   | 26 ++++-----
 torchTextClassifiers/model/lightning.py      | 26 ++++-----
 torchTextClassifiers/torchTextClassifiers.py | 23 ++++----
 5 files changed, 75 insertions(+), 65 deletions(-)

diff --git a/docs/source/architecture/overview.md b/docs/source/architecture/overview.md
index 34aad06..fe0467b 100644
--- a/docs/source/architecture/overview.md
+++ b/docs/source/architecture/overview.md
@@ -611,9 +611,11 @@ This composes naturally with a loss's own per-*class* `weight=` argument (e.g.
 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) can opt in by adding an
-optional `sample_weights` keyword argument to `forward`, e.g.
-`forward(self, outputs, labels, sample_weights=None)` — `MultiLevelCrossEntropyLoss`
+**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
diff --git a/tests/test_sample_weights.py b/tests/test_sample_weights.py
index 9de8550..7a18fea 100644
--- a/tests/test_sample_weights.py
+++ b/tests/test_sample_weights.py
@@ -145,34 +145,57 @@ def test_zero_weight_equivalent_to_excluding_sample(self):
         assert torch.allclose(loss_with_zero, loss_excluded, atol=1e-6)
 
 
-class TestMultiLevelCrossEntropyLossSampleWeights:
-    def test_weighted_matches_manual_computation(self):
+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)
-        weights = torch.tensor([1.0, 0.0, 2.0, 1.0])
 
         loss_fn = MultiLevelCrossEntropyLoss()
-        loss = loss_fn(outputs, labels, sample_weights=weights)
+        assert loss_fn.reduction == "none"
 
-        per_level_losses = []
-        for i, out in enumerate(outputs):
-            per_sample = torch.nn.functional.cross_entropy(out, labels[:, i], reduction="none")
-            per_level_losses.append((per_sample * weights).sum() / weights.sum())
-        expected = sum(per_level_losses) / len(outputs)
+        per_sample = loss_fn(outputs, labels)
+        assert per_sample.shape == (4,)
 
-        assert torch.allclose(loss, expected, atol=1e-6)
+        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_none_sample_weights_matches_unweighted(self):
+    def test_step_applies_sample_weights_across_levels(self):
         torch.manual_seed(0)
-        outputs = [torch.randn(4, 3)]
-        labels = torch.tensor([0, 1, 2, 1]).unsqueeze(1)
+        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])
 
-        loss_fn = MultiLevelCrossEntropyLoss()
-        weighted = loss_fn(outputs, labels)
-        expected = torch.nn.functional.cross_entropy(outputs[0], labels[:, 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)
 
-        assert torch.allclose(weighted, expected, atol=1e-6)
+        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:
diff --git a/torchTextClassifiers/contrib/multilevel.py b/torchTextClassifiers/contrib/multilevel.py
index 9db1752..b640079 100644
--- a/torchTextClassifiers/contrib/multilevel.py
+++ b/torchTextClassifiers/contrib/multilevel.py
@@ -167,37 +167,29 @@ class MultiLevelCrossEntropyLoss(nn.Module):
     def __init__(self, num_classes: Optional[list[int]] = None):
         super().__init__()
         self.num_classes = num_classes
+        self.reduction = "none"
         self.loss_fn = nn.CrossEntropyLoss(reduction="none")
 
-    def forward(
-        self,
-        outputs: list[torch.Tensor],
-        labels: torch.Tensor,
-        sample_weights: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        """Compute the weighted average loss.
+    def forward(self, outputs: list[torch.Tensor], labels: torch.Tensor) -> torch.Tensor:
+        """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
                 by ``MultiLevelTextClassificationModel``.
             labels: Integer label tensor of shape ``(batch, n_levels)``.
                 Column ``i`` contains the ground-truth label for level ``i``.
-            sample_weights: Optional per-sample weight tensor of shape
-                ``(batch,)``. If ``None``, all samples are weighted equally.
 
         Returns:
-            Scalar loss tensor.
+            Per-sample loss tensor of shape ``(batch,)``.
         """
-        if sample_weights is None:
-            sample_weights = torch.ones(outputs[0].shape[0], device=outputs[0].device)
-
-        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
-            per_sample_loss = self.loss_fn(output.squeeze(), label)
-            level_loss = (per_sample_loss * sample_weights).sum() / sample_weights.sum()
-            total_loss = total_loss + level_loss * weight
+            total_loss = total_loss + self.loss_fn(output.squeeze(), label) * weight
 
         total_weight = sum(self.num_classes) if self.num_classes is not None else len(outputs)
         return total_loss / total_weight
diff --git a/torchTextClassifiers/model/lightning.py b/torchTextClassifiers/model/lightning.py
index 5c38e8b..dea98dc 100644
--- a/torchTextClassifiers/model/lightning.py
+++ b/torchTextClassifiers/model/lightning.py
@@ -1,4 +1,3 @@
-import inspect
 import logging
 
 import pytorch_lightning as pl
@@ -45,16 +44,12 @@ def __init__(
         self.model = model
         self.loss = loss
 
-        self._loss_accepts_sample_weights = "sample_weights" in inspect.signature(
-            self.loss.forward
-        ).parameters
-        if not self._loss_accepts_sample_weights and hasattr(self.loss, "reduction"):
-            if self.loss.reduction != "none":
-                logger.info(
-                    f"Setting reduction='none' on {type(self.loss).__name__} so that "
-                    "sample_weights can be applied per-sample before averaging."
-                )
-                self.loss.reduction = "none"
+        if hasattr(self.loss, "reduction") and self.loss.reduction != "none":
+            logger.info(
+                f"Setting reduction='none' on {type(self.loss).__name__} so that "
+                "sample_weights can be applied per-sample before averaging."
+            )
+            self.loss.reduction = "none"
 
         if not hasattr(self.model, "num_classes") or self.model.num_classes is None:
             raise ValueError("Model must have num_classes attribute for accuracy calculation.")
@@ -100,12 +95,9 @@ def step(self, batch) -> tuple[torch.Tensor, torch.Tensor | list[torch.Tensor]]:
             sample_weights = torch.ones(targets.shape[0], device=targets.device)
         sample_weights = sample_weights.to(targets.device)
 
-        if self._loss_accepts_sample_weights:
-            loss = self.loss(outputs, targets, sample_weights=sample_weights)
-        else:
-            per_sample_loss = self.loss(outputs, targets)
-            per_sample_loss = per_sample_loss.reshape(per_sample_loss.size(0), -1).mean(dim=1)
-            loss = (per_sample_loss * sample_weights).sum() / sample_weights.sum()
+        per_sample_loss = self.loss(outputs, targets)
+        per_sample_loss = per_sample_loss.reshape(per_sample_loss.size(0), -1).mean(dim=1)
+        loss = (per_sample_loss * sample_weights).sum() / sample_weights.sum()
 
         if self.multilevel_accuracy:
             accuracy = [
diff --git a/torchTextClassifiers/torchTextClassifiers.py b/torchTextClassifiers/torchTextClassifiers.py
index 3200aef..cc0fdc4 100644
--- a/torchTextClassifiers/torchTextClassifiers.py
+++ b/torchTextClassifiers/torchTextClassifiers.py
@@ -281,17 +281,18 @@ def forward(
            The wrapper reads ``categorical_variable_net.categorical_vocabulary_sizes``
            to set up the data pipeline.
 
-        **Optional: ``sample_weights`` support in a custom loss** — to support
-        per-sample loss weighting (see ``torchTextClassifiers.train``'s
-        ``sample_weights``/``val_sample_weights`` arguments) with a custom
-        multi-task loss, add an optional ``sample_weights`` keyword argument to
-        its ``forward`` method, e.g. ``forward(self, outputs, labels,
-        sample_weights=None)``, and use it to compute a weighted average
-        instead of a plain mean (see ``MultiLevelCrossEntropyLoss`` below for
-        an example). If your loss exposes a ``reduction`` attribute instead
-        (like standard ``torch.nn.*Loss`` classes) it will automatically be
-        switched to ``"none"`` so that the wrapper can apply the weights and
-        reduce the loss itself.
+        **``sample_weights`` support in a custom loss** — per-sample loss
+        weighting (see ``torchTextClassifiers.train``'s
+        ``sample_weights``/``val_sample_weights`` arguments) is applied by
+        ``TextClassificationModule`` itself, outside of your loss. Your loss's
+        ``forward`` just needs to return an **unreduced, per-sample** tensor
+        of shape ``(batch,)`` — exactly like a standard ``torch.nn.*Loss``
+        constructed with ``reduction="none"`` — instead of a single scalar.
+        If your loss exposes a ``reduction`` attribute (like standard
+        ``torch.nn.*Loss`` classes) it is automatically switched to
+        ``"none"``; if not (as for custom multi-task losses), simply make it
+        return the per-sample tensor directly (see
+        ``MultiLevelCrossEntropyLoss`` below for an example).
 
         See ``torchTextClassifiers.contrib`` for ready-made example architectures
         (``MultiLevelTextClassificationModel``, ``MultiLevelCrossEntropyLoss``) that