[kashif]

What does this PR do?

JEV currently returns only logits, so training needs a custom loss outside the model. This adds JevForDecisionLoss through LOSS_MAPPING and returns standard loss/logits outputs for Trainer and TRL, while preserving decision-pipeline inference.

The default follows the AutoTrust JEV-9B model card: KL(target || model) over active options, plus 0.5 × RPS for Score questions. Includes hard/soft targets, question weights, training docs, and loss/gradient tests. The loss uses max-excluded log1p normalization to retain confident-prediction gradients.

Extracted from #50411 against main; that PR stays a draft.

Validation: 20 tests passed, 4 skipped, 31 subtests passed; 128 independent float64 objective/gradient comparisons passed. make style passes. make check-repo only reports the existing torch.distributed.get_rank type-check warning, also reproduced on unchanged main.

Who can review?

Anyone interested in decision-model training.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Pages changed in this PR:

• en/model_doc/jev

[mangeshraut712]

Nice work: the packed KL + optional RPS objective is carefully written, decide() now forces return_dict=True and reads .logits so it keeps working with JevDecisionOutput, and the independent Categorical/RPS reference test (ragged questions, ignore index, soft labels, weights, confident fp32 gradients, microbatch accumulation) is excellent.

A few residual issues:

1. Soft labels are documented to "sum to one" but never checked. JevForDecisionLoss will happily take a soft row that sums to 0.5 or 1.5; kl_div will still run and RPS CDFs will be biased. A cheap targets.sum(-1) check on the active slots (or a test that this is rejected / renormalized) would match the docstring.

2. num_options == 0 / empty packed logits. positions = torch.arange(num_options.max()) and scores.argmax(-1) are undefined or awkward on an empty question. Even a single-question test with zero options (or an empty logits) would show whether this should raise a clear error.

3. if rps_weight: treats a programmed 0.0 correctly, but a negative weight would still be applied. If the weight is user-facing on JevConfig, a < 0 guard might be worth it.

4. Shape of logits. The loss indexes logits as 1-D (logits[indices], logits.numel()). The new output type says shape (total_num_options,), which matches the current forward. If a caller ever passes a batched (1, total) tensor, numel() and advanced indexing will go wrong. Worth asserting logits.ndim == 1 (or flattening a leading batch-1 dim) so a future wrapper cannot fail silently.

accepts_loss_kwargs = False is a good call given packed soft labels, and normalizing by valid-question count even when loss_weights is set is consistent with the new docs.

[ArthurZucker]

NiceQ Smalls nits but will be nice

[ArthurZucker]

maybe return tuple decortor missing?

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: jev