Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / infer / sampler.py: 0%

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"""Sampling helpers for autoregressive generation.""" 

16import torch 

17from torch.nn import functional as F 

18 

19from hyper_parallel.infer.utils import GenerationConfig 

20 

21 

22def _filter_top_k(logits: torch.Tensor, top_k: int) -> torch.Tensor: 

23 if top_k <= 0 or top_k >= logits.size(-1): 

24 return logits 

25 values, _ = torch.topk(logits, k=top_k, dim=-1) 

26 threshold = values[:, -1:].contiguous() 

27 return logits.masked_fill(logits < threshold, float("-inf")) 

28 

29 

30def greedy_sample(logits: torch.Tensor) -> torch.Tensor: 

31 """Select the highest-logit token for each batch item.""" 

32 if logits.ndim != 2: 

33 raise ValueError("logits must have shape (batch, vocab)") 

34 return logits.argmax(dim=-1, keepdim=True) 

35 

36 

37def top_k_sample( 

38 logits: torch.Tensor, 

39 top_k: int, 

40 temperature: float = 1.0, 

41) -> torch.Tensor: 

42 """Sample from the top-k logits.""" 

43 if logits.ndim != 2: 

44 raise ValueError("logits must have shape (batch, vocab)") 

45 filtered = _filter_top_k(logits, top_k) 

46 probs = F.softmax(filtered / temperature, dim=-1) 

47 sampled = torch.multinomial(probs, num_samples=1) 

48 return sampled 

49 

50 

51def top_p_sample( 

52 logits: torch.Tensor, 

53 top_p: float, 

54 temperature: float = 1.0, 

55) -> torch.Tensor: 

56 """Sample from the nucleus token set.""" 

57 if logits.ndim != 2: 

58 raise ValueError("logits must have shape (batch, vocab)") 

59 if top_p >= 1.0: 

60 probs = F.softmax(logits / temperature, dim=-1) 

61 return torch.multinomial(probs, num_samples=1) 

62 sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1) 

63 sorted_probs = F.softmax(sorted_logits / temperature, dim=-1) 

64 cumulative = sorted_probs.cumsum(dim=-1) 

65 remove = cumulative - sorted_probs > top_p 

66 filtered = sorted_logits.masked_fill(remove, float("-inf")) 

67 probs = F.softmax(filtered / temperature, dim=-1) 

68 sampled = torch.multinomial(probs, num_samples=1) 

69 return sorted_indices.gather(dim=-1, index=sampled) 

70 

71 

72def apply_repetition_penalty( 

73 logits: torch.Tensor, 

74 input_ids: torch.Tensor, 

75 penalty: float, 

76) -> torch.Tensor: 

77 """Apply per-item repetition penalty to seen token ids.""" 

78 if penalty == 1.0: 

79 return logits 

80 if logits.ndim != 2 or input_ids.ndim != 2: 

81 raise ValueError("logits and input_ids must be 2-D tensors") 

82 adjusted = logits.clone() 

83 vocab_size = logits.size(-1) 

84 valid = (input_ids >= 0) & (input_ids < vocab_size) 

85 seen_mask = torch.zeros_like(adjusted, dtype=torch.bool) 

86 token_ids = input_ids.to(dtype=torch.long).clamp(min=0, max=max(vocab_size - 1, 0)) 

87 seen_mask.scatter_(dim=1, index=token_ids, src=valid) 

88 penalized = torch.where(adjusted < 0, adjusted * penalty, adjusted / penalty) 

89 return torch.where(seen_mask, penalized, adjusted) 

90 

91 

92def sample_next_token( 

93 logits: torch.Tensor, 

94 input_ids: torch.Tensor, 

95 config: GenerationConfig, 

96) -> torch.Tensor: 

97 """Apply repetition penalty and select the next token.""" 

98 logits = apply_repetition_penalty( 

99 logits, 

100 input_ids=input_ids, 

101 penalty=config.repetition_penalty, 

102 ) 

103 if not config.do_sample: 

104 return greedy_sample(logits) 

105 logits = _filter_top_k(logits, config.top_k) 

106 if config.top_p < 1.0: 

107 return top_p_sample(logits, top_p=config.top_p, temperature=config.temperature) 

108 probs = F.softmax(logits / config.temperature, dim=-1) 

109 return torch.multinomial(probs, num_samples=1)