add file
Browse files- qwen_generation_utils.py +416 -0
qwen_generation_utils.py
ADDED
@@ -0,0 +1,416 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Alibaba Cloud.
|
2 |
+
#
|
3 |
+
# This source code is licensed under the license found in the
|
4 |
+
# LICENSE file in the root directory of this source tree.
|
5 |
+
|
6 |
+
"""Generation support."""
|
7 |
+
|
8 |
+
from typing import Tuple, List, Union, Iterable
|
9 |
+
|
10 |
+
import numpy as np
|
11 |
+
import torch
|
12 |
+
import torch.nn.functional as F
|
13 |
+
from transformers import PreTrainedTokenizer
|
14 |
+
from transformers import logging
|
15 |
+
from transformers.generation import LogitsProcessor
|
16 |
+
|
17 |
+
logger = logging.get_logger(__name__)
|
18 |
+
|
19 |
+
# Types.
|
20 |
+
HistoryType = List[Tuple[str, str]]
|
21 |
+
TokensType = List[int]
|
22 |
+
BatchTokensType = List[List[int]]
|
23 |
+
|
24 |
+
|
25 |
+
def pad_batch(batch: BatchTokensType, pad_id: int, seq_length: int) -> BatchTokensType:
|
26 |
+
for tokens in batch:
|
27 |
+
context_length = len(tokens)
|
28 |
+
if context_length < seq_length:
|
29 |
+
tokens.extend([pad_id] * (seq_length - context_length))
|
30 |
+
return batch
|
31 |
+
|
32 |
+
|
33 |
+
def get_ltor_masks_and_position_ids(
|
34 |
+
data,
|
35 |
+
eod_token,
|
36 |
+
reset_position_ids,
|
37 |
+
reset_attention_mask,
|
38 |
+
eod_mask_loss,
|
39 |
+
):
|
40 |
+
"""Build masks and position id for left to right model."""
|
41 |
+
|
42 |
+
# Extract batch size and sequence length.
|
43 |
+
micro_batch_size, seq_length = data.size()
|
44 |
+
|
45 |
+
# Attention mask (lower triangular).
|
46 |
+
if reset_attention_mask:
|
47 |
+
att_mask_batch = micro_batch_size
|
48 |
+
else:
|
49 |
+
att_mask_batch = 1
|
50 |
+
attention_mask = torch.tril(
|
51 |
+
torch.ones((att_mask_batch, seq_length, seq_length), device=data.device)
|
52 |
+
).view(att_mask_batch, 1, seq_length, seq_length)
|
53 |
+
|
54 |
+
# Loss mask.
|
55 |
+
loss_mask = torch.ones(data.size(), dtype=torch.float, device=data.device)
|
56 |
+
if eod_mask_loss:
|
57 |
+
loss_mask[data == eod_token] = 0.0
|
58 |
+
|
59 |
+
# Position ids.
|
60 |
+
position_ids = torch.arange(seq_length, dtype=torch.long, device=data.device)
|
61 |
+
position_ids = position_ids.unsqueeze(0).expand_as(data)
|
62 |
+
# We need to clone as the ids will be modifed based on batch index.
|
63 |
+
if reset_position_ids:
|
64 |
+
position_ids = position_ids.clone()
|
65 |
+
|
66 |
+
if reset_position_ids or reset_attention_mask:
|
67 |
+
# Loop through the batches:
|
68 |
+
for b in range(micro_batch_size):
|
69 |
+
|
70 |
+
# Find indecies where EOD token is.
|
71 |
+
eod_index = position_ids[b, data[b] == eod_token]
|
72 |
+
# Detach indecies from positions if going to modify positions.
|
73 |
+
if reset_position_ids:
|
74 |
+
eod_index = eod_index.clone()
|
75 |
+
|
76 |
+
# Loop through EOD indecies:
|
77 |
+
prev_index = 0
|
78 |
+
for j in range(eod_index.size()[0]):
|
79 |
+
i = eod_index[j]
|
80 |
+
# Mask attention loss.
|
81 |
+
if reset_attention_mask:
|
82 |
+
attention_mask[b, 0, (i + 1) :, : (i + 1)] = 0
|
83 |
+
# Reset positions.
|
84 |
+
if reset_position_ids:
|
85 |
+
position_ids[b, (i + 1) :] -= i + 1 - prev_index
|
86 |
+
prev_index = i + 1
|
87 |
+
|
88 |
+
# Convert attention mask to binary:
|
89 |
+
attention_mask = attention_mask < 0.5
|
90 |
+
|
91 |
+
return attention_mask, loss_mask, position_ids
|
92 |
+
|
93 |
+
|
94 |
+
def get_batch(context_tokens: torch.LongTensor, eod_id: int):
|
95 |
+
"""Generate batch from context tokens."""
|
96 |
+
# Move to GPU.
|
97 |
+
tokens = context_tokens.contiguous().to(context_tokens.device)
|
98 |
+
# Get the attention mask and postition ids.
|
99 |
+
attention_mask, _, position_ids = get_ltor_masks_and_position_ids(
|
100 |
+
tokens,
|
101 |
+
eod_id,
|
102 |
+
reset_position_ids=False,
|
103 |
+
reset_attention_mask=False,
|
104 |
+
eod_mask_loss=False,
|
105 |
+
)
|
106 |
+
return tokens, attention_mask, position_ids
|
107 |
+
|
108 |
+
|
109 |
+
def get_stop_words_ids(chat_format, tokenizer):
|
110 |
+
if chat_format == "raw":
|
111 |
+
stop_words_ids = [tokenizer.encode("Human:"), [tokenizer.eod_id]]
|
112 |
+
elif chat_format == "chatml":
|
113 |
+
stop_words_ids = [[tokenizer.im_end_id], [tokenizer.im_start_id]]
|
114 |
+
else:
|
115 |
+
raise NotImplementedError(f"Unknown chat format {chat_format!r}")
|
116 |
+
return stop_words_ids
|
117 |
+
|
118 |
+
|
119 |
+
def make_context(
|
120 |
+
tokenizer: PreTrainedTokenizer,
|
121 |
+
query: str,
|
122 |
+
history: List[Tuple[str, str]] = None,
|
123 |
+
system: str = "",
|
124 |
+
max_window_size: int = 6144,
|
125 |
+
chat_format: str = "chatml",
|
126 |
+
):
|
127 |
+
if history is None:
|
128 |
+
history = []
|
129 |
+
|
130 |
+
if chat_format == "chatml":
|
131 |
+
im_start, im_end = "<|im_start|>", "<|im_end|>"
|
132 |
+
im_start_tokens = [tokenizer.im_start_id]
|
133 |
+
im_end_tokens = [tokenizer.im_end_id]
|
134 |
+
nl_tokens = tokenizer.encode("\n")
|
135 |
+
|
136 |
+
def _tokenize_str(role, content):
|
137 |
+
return f"{role}\n{content}", tokenizer.encode(
|
138 |
+
role, allowed_special=set()
|
139 |
+
) + nl_tokens + tokenizer.encode(content, allowed_special=set())
|
140 |
+
|
141 |
+
system_text, system_tokens_part = _tokenize_str("system", system)
|
142 |
+
system_tokens = im_start_tokens + system_tokens_part + im_end_tokens
|
143 |
+
|
144 |
+
raw_text = ""
|
145 |
+
context_tokens = []
|
146 |
+
|
147 |
+
for turn_query, turn_response in reversed(history):
|
148 |
+
query_text, query_tokens_part = _tokenize_str("user", turn_query)
|
149 |
+
query_tokens = im_start_tokens + query_tokens_part + im_end_tokens
|
150 |
+
response_text, response_tokens_part = _tokenize_str(
|
151 |
+
"assistant", turn_response
|
152 |
+
)
|
153 |
+
response_tokens = im_start_tokens + response_tokens_part + im_end_tokens
|
154 |
+
|
155 |
+
next_context_tokens = nl_tokens + query_tokens + nl_tokens + response_tokens
|
156 |
+
prev_chat = (
|
157 |
+
f"\n{im_start}{query_text}{im_end}\n{im_start}{response_text}{im_end}"
|
158 |
+
)
|
159 |
+
|
160 |
+
current_context_size = (
|
161 |
+
len(system_tokens) + len(next_context_tokens) + len(context_tokens)
|
162 |
+
)
|
163 |
+
if current_context_size < max_window_size:
|
164 |
+
context_tokens = next_context_tokens + context_tokens
|
165 |
+
raw_text = prev_chat + raw_text
|
166 |
+
else:
|
167 |
+
break
|
168 |
+
|
169 |
+
context_tokens = system_tokens + context_tokens
|
170 |
+
raw_text = f"{im_start}{system_text}{im_end}" + raw_text
|
171 |
+
context_tokens += (
|
172 |
+
nl_tokens
|
173 |
+
+ im_start_tokens
|
174 |
+
+ _tokenize_str("user", query)[1]
|
175 |
+
+ im_end_tokens
|
176 |
+
+ nl_tokens
|
177 |
+
+ im_start_tokens
|
178 |
+
+ tokenizer.encode("assistant")
|
179 |
+
+ nl_tokens
|
180 |
+
)
|
181 |
+
raw_text += f"\n{im_start}user\n{query}{im_end}\n{im_start}assistant\n"
|
182 |
+
|
183 |
+
elif chat_format == "raw":
|
184 |
+
raw_text = query
|
185 |
+
context_tokens = tokenizer.encode(raw_text)
|
186 |
+
else:
|
187 |
+
raise NotImplementedError(f"Unknown chat format {chat_format!r}")
|
188 |
+
|
189 |
+
return raw_text, context_tokens
|
190 |
+
|
191 |
+
|
192 |
+
def _decode_default(
|
193 |
+
tokens: List[int],
|
194 |
+
*,
|
195 |
+
stop_words: List[str],
|
196 |
+
eod_words: List[str],
|
197 |
+
tokenizer: PreTrainedTokenizer,
|
198 |
+
raw_text_len: int,
|
199 |
+
verbose: bool = False,
|
200 |
+
return_end_reason: bool = False,
|
201 |
+
errors: str='replace',
|
202 |
+
):
|
203 |
+
trim_decode_tokens = tokenizer.decode(tokens, errors=errors)[raw_text_len:]
|
204 |
+
if verbose:
|
205 |
+
print("\nRaw Generate: ", trim_decode_tokens)
|
206 |
+
|
207 |
+
end_reason = f"Gen length {len(tokens)}"
|
208 |
+
for stop_word in stop_words:
|
209 |
+
trim_decode_tokens = trim_decode_tokens.replace(stop_word, "").strip()
|
210 |
+
for eod_word in eod_words:
|
211 |
+
if eod_word in trim_decode_tokens:
|
212 |
+
end_reason = f"Gen {eod_word!r}"
|
213 |
+
trim_decode_tokens = trim_decode_tokens.split(eod_word)[0]
|
214 |
+
trim_decode_tokens = trim_decode_tokens.strip()
|
215 |
+
if verbose:
|
216 |
+
print("\nEnd Reason:", end_reason)
|
217 |
+
print("\nGenerate: ", trim_decode_tokens)
|
218 |
+
|
219 |
+
if return_end_reason:
|
220 |
+
return trim_decode_tokens, end_reason
|
221 |
+
else:
|
222 |
+
return trim_decode_tokens
|
223 |
+
|
224 |
+
|
225 |
+
def _decode_chatml(
|
226 |
+
tokens: List[int],
|
227 |
+
*,
|
228 |
+
stop_words: List[str],
|
229 |
+
eod_token_ids: List[int],
|
230 |
+
tokenizer: PreTrainedTokenizer,
|
231 |
+
raw_text_len: int,
|
232 |
+
context_length: int,
|
233 |
+
verbose: bool = False,
|
234 |
+
return_end_reason: bool = False,
|
235 |
+
errors: str='replace'
|
236 |
+
):
|
237 |
+
end_reason = f"Gen length {len(tokens)}"
|
238 |
+
eod_token_idx = context_length
|
239 |
+
for eod_token_idx in range(context_length, len(tokens)):
|
240 |
+
if tokens[eod_token_idx] in eod_token_ids:
|
241 |
+
end_reason = f"Gen {tokenizer.decode([tokens[eod_token_idx]])!r}"
|
242 |
+
break
|
243 |
+
|
244 |
+
trim_decode_tokens = tokenizer.decode(tokens[:eod_token_idx], errors=errors)[raw_text_len:]
|
245 |
+
if verbose:
|
246 |
+
print("\nRaw Generate w/o EOD:", tokenizer.decode(tokens, errors=errors)[raw_text_len:])
|
247 |
+
print("\nRaw Generate:", trim_decode_tokens)
|
248 |
+
print("\nEnd Reason:", end_reason)
|
249 |
+
for stop_word in stop_words:
|
250 |
+
trim_decode_tokens = trim_decode_tokens.replace(stop_word, "").strip()
|
251 |
+
trim_decode_tokens = trim_decode_tokens.strip()
|
252 |
+
if verbose:
|
253 |
+
print("\nGenerate:", trim_decode_tokens)
|
254 |
+
|
255 |
+
if return_end_reason:
|
256 |
+
return trim_decode_tokens, end_reason
|
257 |
+
else:
|
258 |
+
return trim_decode_tokens
|
259 |
+
|
260 |
+
|
261 |
+
def decode_tokens(
|
262 |
+
tokens: Union[torch.LongTensor, TokensType],
|
263 |
+
tokenizer: PreTrainedTokenizer,
|
264 |
+
raw_text_len: int,
|
265 |
+
context_length: int,
|
266 |
+
chat_format: str,
|
267 |
+
verbose: bool = False,
|
268 |
+
return_end_reason: bool = False,
|
269 |
+
errors: str="replace",
|
270 |
+
) -> str:
|
271 |
+
if torch.is_tensor(tokens):
|
272 |
+
tokens = tokens.cpu().numpy().tolist()
|
273 |
+
|
274 |
+
if chat_format == "chatml":
|
275 |
+
return _decode_chatml(
|
276 |
+
tokens,
|
277 |
+
stop_words=[],
|
278 |
+
eod_token_ids=[tokenizer.im_start_id, tokenizer.im_end_id],
|
279 |
+
tokenizer=tokenizer,
|
280 |
+
raw_text_len=raw_text_len,
|
281 |
+
context_length=context_length,
|
282 |
+
verbose=verbose,
|
283 |
+
return_end_reason=return_end_reason,
|
284 |
+
errors=errors,
|
285 |
+
)
|
286 |
+
elif chat_format == "raw":
|
287 |
+
return _decode_default(
|
288 |
+
tokens,
|
289 |
+
stop_words=["<|endoftext|>"],
|
290 |
+
eod_words=["<|endoftext|>"],
|
291 |
+
tokenizer=tokenizer,
|
292 |
+
raw_text_len=raw_text_len,
|
293 |
+
verbose=verbose,
|
294 |
+
return_end_reason=return_end_reason,
|
295 |
+
errors=errors,
|
296 |
+
)
|
297 |
+
else:
|
298 |
+
raise NotImplementedError(f"Unknown chat format {chat_format!r}")
|
299 |
+
|
300 |
+
|
301 |
+
class StopWordsLogitsProcessor(LogitsProcessor):
|
302 |
+
"""
|
303 |
+
:class:`transformers.LogitsProcessor` that enforces that when specified sequences appear, stop geration.
|
304 |
+
|
305 |
+
Args:
|
306 |
+
stop_words_ids (:obj:`List[List[int]]`):
|
307 |
+
List of list of token ids of stop ids. In order to get the tokens of the words
|
308 |
+
that should not appear in the generated text, use :obj:`tokenizer(bad_word,
|
309 |
+
add_prefix_space=True).input_ids`.
|
310 |
+
eos_token_id (:obj:`int`):
|
311 |
+
The id of the `end-of-sequence` token.
|
312 |
+
"""
|
313 |
+
|
314 |
+
def __init__(self, stop_words_ids: Iterable[Iterable[int]], eos_token_id: int):
|
315 |
+
|
316 |
+
if not isinstance(stop_words_ids, List) or len(stop_words_ids) == 0:
|
317 |
+
raise ValueError(
|
318 |
+
f"`stop_words_ids` has to be a non-emtpy list, but is {stop_words_ids}."
|
319 |
+
)
|
320 |
+
if any(not isinstance(bad_word_ids, list) for bad_word_ids in stop_words_ids):
|
321 |
+
raise ValueError(
|
322 |
+
f"`stop_words_ids` has to be a list of lists, but is {stop_words_ids}."
|
323 |
+
)
|
324 |
+
if any(
|
325 |
+
any(
|
326 |
+
(not isinstance(token_id, (int, np.integer)) or token_id < 0)
|
327 |
+
for token_id in stop_word_ids
|
328 |
+
)
|
329 |
+
for stop_word_ids in stop_words_ids
|
330 |
+
):
|
331 |
+
raise ValueError(
|
332 |
+
f"Each list in `stop_words_ids` has to be a list of positive integers, but is {stop_words_ids}."
|
333 |
+
)
|
334 |
+
|
335 |
+
self.stop_words_ids = list(
|
336 |
+
filter(
|
337 |
+
lambda bad_token_seq: bad_token_seq != [eos_token_id], stop_words_ids
|
338 |
+
)
|
339 |
+
)
|
340 |
+
self.eos_token_id = eos_token_id
|
341 |
+
for stop_token_seq in self.stop_words_ids:
|
342 |
+
assert (
|
343 |
+
len(stop_token_seq) > 0
|
344 |
+
), "Stop words token sequences {} cannot have an empty list".format(
|
345 |
+
stop_words_ids
|
346 |
+
)
|
347 |
+
|
348 |
+
def __call__(
|
349 |
+
self, input_ids: torch.LongTensor, scores: torch.FloatTensor
|
350 |
+
) -> torch.FloatTensor:
|
351 |
+
stopped_samples = self._calc_stopped_samples(input_ids)
|
352 |
+
for i, should_stop in enumerate(stopped_samples):
|
353 |
+
if should_stop:
|
354 |
+
scores[i, self.eos_token_id] = float(2**15)
|
355 |
+
return scores
|
356 |
+
|
357 |
+
def _tokens_match(self, prev_tokens: torch.LongTensor, tokens: List[int]) -> bool:
|
358 |
+
if len(tokens) == 0:
|
359 |
+
# if bad word tokens is just one token always ban it
|
360 |
+
return True
|
361 |
+
elif len(tokens) > len(prev_tokens):
|
362 |
+
# if bad word tokens are longer then prev input_ids they can't be equal
|
363 |
+
return False
|
364 |
+
elif prev_tokens[-len(tokens) :].tolist() == tokens:
|
365 |
+
# if tokens match
|
366 |
+
return True
|
367 |
+
else:
|
368 |
+
return False
|
369 |
+
|
370 |
+
def _calc_stopped_samples(self, prev_input_ids: Iterable[int]) -> Iterable[int]:
|
371 |
+
stopped_samples = []
|
372 |
+
for prev_input_ids_slice in prev_input_ids:
|
373 |
+
match = False
|
374 |
+
for stop_token_seq in self.stop_words_ids:
|
375 |
+
if self._tokens_match(prev_input_ids_slice, stop_token_seq):
|
376 |
+
# if tokens do not match continue
|
377 |
+
match = True
|
378 |
+
break
|
379 |
+
stopped_samples.append(match)
|
380 |
+
|
381 |
+
return stopped_samples
|
382 |
+
|
383 |
+
|
384 |
+
def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-float("Inf")):
|
385 |
+
"""This function has been mostly taken from huggingface conversational
|
386 |
+
ai code at
|
387 |
+
https://medium.com/huggingface/how-to-build-a-state-of-the-art-
|
388 |
+
conversational-ai-with-transfer-learning-2d818ac26313"""
|
389 |
+
|
390 |
+
if top_k > 0:
|
391 |
+
# Remove all tokens with a probability less than the
|
392 |
+
# last token of the top-k
|
393 |
+
indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
|
394 |
+
logits[indices_to_remove] = filter_value
|
395 |
+
|
396 |
+
if top_p > 0.0:
|
397 |
+
# Cconvert to 1D
|
398 |
+
sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
|
399 |
+
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
|
400 |
+
|
401 |
+
# Remove tokens with cumulative probability above the threshold
|
402 |
+
sorted_indices_to_remove = cumulative_probs > top_p
|
403 |
+
# Shift the indices to the right to keep also the first token
|
404 |
+
# above the threshold
|
405 |
+
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
|
406 |
+
sorted_indices_to_remove[..., 0] = 0
|
407 |
+
for i in range(sorted_indices.size(0)):
|
408 |
+
indices_to_remove = sorted_indices[i][sorted_indices_to_remove[i]]
|
409 |
+
logits[i][indices_to_remove] = filter_value
|
410 |
+
|
411 |
+
return logits
|
412 |
+
|
413 |
+
|
414 |
+
def switch(val1, val2, boolean):
|
415 |
+
boolean = boolean.type_as(val1)
|
416 |
+
return (1 - boolean) * val1 + boolean * val2
|