Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / losses / masked_ce.py: 37%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Directly configurable masked cross-entropy callable.""" 

16 

17__all__ = ["MaskedCrossEntropy"] 

18 

19from dataclasses import dataclass 

20from typing import Literal, Optional 

21 

22import torch 

23import torch.nn.functional as F 

24 

25 

26@dataclass(kw_only=True, slots=True) # pylint: disable=unexpected-keyword-arg 

27class MaskedCrossEntropy: 

28 """YAML-targeted masked cross-entropy callable.""" 

29 

30 fp32_upcast: bool = True 

31 ignore_index: int = -100 

32 reduction: Literal["none", "mean", "sum"] = "sum" 

33 

34 def __call__( 

35 self, 

36 logits: torch.Tensor, 

37 labels: torch.Tensor, 

38 mask: Optional[torch.Tensor] = None, 

39 num_label_tokens: Optional[int] = None, 

40 ) -> torch.Tensor: 

41 """Compute the configured loss. 

42 

43 Args: 

44 logits: Model logits whose final dimension is the vocabulary. 

45 labels: Target token IDs. 

46 mask: Optional mask where zero marks ignored target positions. 

47 num_label_tokens: Optional normalization denominator. 

48 

49 Returns: 

50 Masked cross-entropy loss. 

51 

52 Raises: 

53 ValueError: If token normalization is requested for a non-sum loss. 

54 """ 

55 if labels.device != logits.device: 

56 labels = labels.to(logits.device) 

57 

58 logits = logits.view(-1, logits.size(-1)) 

59 labels = labels.view(-1) 

60 

61 if mask is not None: 

62 with torch.no_grad(): 

63 if mask.device != labels.device: 

64 mask = mask.to(labels.device) 

65 labels = labels.masked_fill(mask.view(-1) == 0, self.ignore_index) 

66 

67 if self.fp32_upcast: 

68 logits = logits.float() 

69 

70 loss = F.cross_entropy( 

71 logits, 

72 labels, 

73 reduction=self.reduction, 

74 ignore_index=self.ignore_index, 

75 ) 

76 

77 if num_label_tokens is not None: 

78 if self.reduction != "sum": 

79 raise ValueError( 

80 "num_label_tokens is only supported when reduction is 'sum'" 

81 ) 

82 if num_label_tokens == 0: 

83 return loss * 0.0 

84 loss = loss / num_label_tokens 

85 

86 return loss