Aluode/PerceptionLabPortable
0
1from abc import ABC, abstractmethod2from typing import Optional3 4from ..utils import logging5 6 7logger = logging.get_logger(__name__)8 9 10# TODO joao, manuel: remove in v4.58.011class Constraint(ABC):12 r"""Abstract base class for all constraints that can be applied during generation.13 It must define how the constraint can be satisfied.14 15 All classes that inherit Constraint must follow the requirement that16 17 ```py18 completed = False19 while not completed:20 _, completed = constraint.update(constraint.advance())21 ```22 23 will always terminate (halt).24 """25 26 def __init__(self):27 logger.warning_once(28 "Importing `Constraint` classes is deprecated and will be removed in v4.58.0. Constrained beam search has been moved to the Hub: https://hf.co/transformers-community/constrained-beam-search. Please import using `from transformers.generation import Constraint` instead."29 )30 # test for the above condition31 self.test()32 33 def test(self):34 """35 Tests whether this constraint has been properly defined.36 """37 counter = 038 completed = False39 while not completed:40 if counter == 1:41 self.reset()42 advance = self.advance()43 if not self.does_advance(advance):44 raise Exception(45 "Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true."46 )47 48 stepped, completed, reset = self.update(advance)49 counter += 150 51 if counter > 10000:52 raise Exception("update() does not fulfill the constraint.")53 54 if self.remaining() != 0:55 raise Exception("Custom Constraint is not defined correctly.")56 57 @abstractmethod58 def advance(self):59 """60 When called, returns the token(s) that would take this constraint one step closer to being fulfilled.61 62 Return:63 token_ids (Union[int, list[int], None]):64 - A single token ID (int) that advances the constraint, or65 - A list of token IDs that could advance the constraint66 - None if the constraint is completed or cannot be advanced67 """68 raise NotImplementedError(69 f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."70 )71 72 @abstractmethod73 def does_advance(self, token_id: int):74 """75 Reads in a token and returns whether it creates progress.76 """77 raise NotImplementedError(78 f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."79 )80 81 @abstractmethod82 def update(self, token_id: int):83 """84 Reads in a token and returns booleans that indicate the progress made by it. This function will update the85 state of this object unlikes `does_advance(self, token_id: int)`.86 87 This isn't to test whether a certain token will advance the progress; it's to update its state as if it has88 been generated. This becomes important if token_id != desired token (refer to else statement in89 PhrasalConstraint)90 91 Args:92 token_id(`int`):93 The id of a newly generated token in the beam search.94 Return:95 stepped(`bool`):96 Whether this constraint has become one step closer to being fulfuilled.97 completed(`bool`):98 Whether this constraint has been completely fulfilled by this token being generated.99 reset (`bool`):100 Whether this constraint has reset its progress by this token being generated.101 """102 raise NotImplementedError(103 f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."104 )105 106 @abstractmethod107 def reset(self):108 """109 Resets the state of this constraint to its initialization. We would call this in cases where the fulfillment of110 a constraint is abrupted by an unwanted token.111 """112 raise NotImplementedError(113 f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."114 )115 116 @abstractmethod117 def remaining(self):118 """119 Returns the number of remaining steps of `advance()` in order to complete this constraint.120 """121 raise NotImplementedError(122 f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."123 )124 125 @abstractmethod126 def copy(self, stateful=False):127 """128 Creates a new instance of this constraint.129 130 Args:131 stateful(`bool`): Whether to not only copy the constraint for new instance, but also its state.132 133 Return:134 constraint(`Constraint`): The same constraint as the one being called from.135 """136 raise NotImplementedError(137 f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."138 )139 140 141class PhrasalConstraint(Constraint):142 r"""143 [`Constraint`] enforcing that an ordered sequence of tokens is included in the output.144 145 Args:146 token_ids (`list[int]`):147 The id of the token that must be generated by the output.148 """149 150 def __init__(self, token_ids: list[int]):151 super(Constraint, self).__init__()152 153 if not isinstance(token_ids, list) or len(token_ids) == 0:154 raise ValueError(f"`token_ids` has to be a non-empty list, but is {token_ids}.")155 if any((not isinstance(token_id, int) or token_id < 0) for token_id in token_ids):156 raise ValueError(f"Each list in `token_ids` has to be a list of positive integers, but is {token_ids}.")157 158 self.token_ids = token_ids159 160 self.seqlen = len(self.token_ids)161 self.fulfilled_idx = -1 # the index of the currently fulfilled step162 self.completed = False163 164 def advance(self):165 if self.completed:166 return None167 return self.token_ids[self.fulfilled_idx + 1]168 169 def does_advance(self, token_id: int):170 if not isinstance(token_id, int):171 raise TypeError(f"`token_id` has to be an `int`, but is {token_id} of type {type(token_id)}")172 173 if self.completed:174 return False175 176 return token_id == self.token_ids[self.fulfilled_idx + 1]177 178 def update(self, token_id: int):179 if not isinstance(token_id, int):180 raise TypeError(f"`token_id` has to be an `int`, but is {token_id} of type {type(token_id)}")181 182 stepped = False183 completed = False184 reset = False185 186 if self.does_advance(token_id):187 self.fulfilled_idx += 1188 stepped = True189 if self.fulfilled_idx == (self.seqlen - 1):190 completed = True191 self.completed = completed192 else:193 # failed to make progress.194 reset = True195 self.reset()196 return stepped, completed, reset197 198 def reset(self):199 self.completed = False200 self.fulfilled_idx = 0201 202 def remaining(self):203 return self.seqlen - (self.fulfilled_idx + 1)204 205 def copy(self, stateful=False):206 new_constraint = PhrasalConstraint(self.token_ids)207 208 if stateful:209 new_constraint.seq_len = self.seqlen210 new_constraint.fulfilled_idx = self.fulfilled_idx211 new_constraint.completed = self.completed212 213 return new_constraint214 215 216class DisjunctiveTrie:217 def __init__(self, nested_token_ids: list[list[int]], no_subsets=True):218 r"""219 A helper class that builds a trie with the words represented in `nested_token_ids`.220 """221 self.max_height = max([len(one) for one in nested_token_ids])222 223 root = {}224 for token_ids in nested_token_ids:225 level = root226 for tidx, token_id in enumerate(token_ids):227 if token_id not in level:228 level[token_id] = {}229 230 level = level[token_id]231 232 if no_subsets and self.has_subsets(root, nested_token_ids):233 raise ValueError(234 "Each list in `nested_token_ids` can't be a complete subset of another list, but is"235 f" {nested_token_ids}."236 )237 238 self.trie = root239 240 def next_tokens(self, current_seq):241 """242 The next possible tokens that will progress the trie, given the current sequence of tokens in `current_seq`.243 """244 start = self.trie245 246 for current_token in current_seq:247 start = start[current_token]248 249 next_tokens = list(start.keys())250 251 return next_tokens252 253 def reached_leaf(self, current_seq):254 next_tokens = self.next_tokens(current_seq)255 256 return len(next_tokens) == 0257 258 def count_leaves(self, root):259 next_nodes = list(root.values())260 if len(next_nodes) == 0:261 return 1262 else:263 return sum([self.count_leaves(nn) for nn in next_nodes])264 265 def has_subsets(self, trie, nested_token_ids):266 """267 Returns whether # of leaves == # of words. Otherwise some word is a subset of another.268 """269 leaf_count = self.count_leaves(trie)270 return len(nested_token_ids) != leaf_count271 272 273class DisjunctiveConstraint(Constraint):274 r"""275 A special [`Constraint`] that is fulfilled by fulfilling just one of several constraints.276 277 Args:278 nested_token_ids (`list[list[int]]`):279 A list of words, where each word is a list of ids. This constraint is fulfilled by generating just one from280 the list of words.281 """282 283 def __init__(self, nested_token_ids: list[list[int]]):284 super(Constraint, self).__init__()285 286 if not isinstance(nested_token_ids, list) or len(nested_token_ids) == 0:287 raise ValueError(f"`nested_token_ids` has to be a non-empty list, but is {nested_token_ids}.")288 if any(not isinstance(token_ids, list) for token_ids in nested_token_ids):289 raise ValueError(f"`nested_token_ids` has to be a list of lists, but is {nested_token_ids}.")290 if any(291 any((not isinstance(token_id, int) or token_id < 0) for token_id in token_ids)292 for token_ids in nested_token_ids293 ):294 raise ValueError(295 f"Each list in `nested_token_ids` has to be a list of positive integers, but is {nested_token_ids}."296 )297 298 self.trie = DisjunctiveTrie(nested_token_ids)299 self.token_ids = nested_token_ids300 301 self.seqlen = self.trie.max_height302 self.current_seq = []303 self.completed = False304 305 def advance(self):306 token_list = self.trie.next_tokens(self.current_seq)307 308 if len(token_list) == 0:309 return None310 else:311 return token_list312 313 def does_advance(self, token_id: int):314 if not isinstance(token_id, int):315 raise TypeError(f"`token_id` is supposed to be type `int`, but is {token_id} of type {type(token_id)}")316 317 next_tokens = self.trie.next_tokens(self.current_seq)318 319 return token_id in next_tokens320 321 def update(self, token_id: int):322 if not isinstance(token_id, int):323 raise TypeError(f"`token_id` is supposed to be type `int`, but is {token_id} of type {type(token_id)}")324 325 stepped = False326 completed = False327 reset = False328 329 if self.does_advance(token_id):330 self.current_seq.append(token_id)331 stepped = True332 else:333 reset = True334 self.reset()335 336 completed = self.trie.reached_leaf(self.current_seq)337 self.completed = completed338 339 return stepped, completed, reset340 341 def reset(self):342 self.completed = False343 self.current_seq = []344 345 def remaining(self):346 if self.completed:347 # since this can be completed without reaching max height348 return 0349 else:350 return self.seqlen - len(self.current_seq)351 352 def copy(self, stateful=False):353 new_constraint = DisjunctiveConstraint(self.token_ids)354 355 if stateful:356 new_constraint.seq_len = self.seqlen357 new_constraint.current_seq = self.current_seq358 new_constraint.completed = self.completed359 360 return new_constraint361 362 363class ConstraintListState:364 r"""365 A class for beam scorers to track its progress through a list of constraints.366 367 Args:368 constraints (`list[Constraint]`):369 A list of [`Constraint`] objects that must be fulfilled by the beam scorer.370 """371 372 def __init__(self, constraints: list[Constraint]):373 self.constraints = constraints374 375 # max # of steps required to fulfill a given constraint376 self.max_seqlen = max([c.seqlen for c in constraints])377 self.n_constraints = len(constraints)378 self.completed = False379 380 self.init_state()381 382 def init_state(self):383 self.complete_constraints = []384 self.inprogress_constraint = None385 self.pending_constraints = [constraint.copy(stateful=False) for constraint in self.constraints]386 387 def get_bank(self):388 add = 0389 if self.inprogress_constraint:390 # extra points for having a constraint mid-fulfilled391 add += self.max_seqlen - self.inprogress_constraint.remaining()392 393 return (len(self.complete_constraints) * self.max_seqlen) + add394 395 def advance(self):396 """The list of tokens to generate such that we can make progress.397 By "list" we don't mean the list of token that will fully fulfill a constraint.398 399 Given constraints `c_i = {t_ij | j == # of tokens}`, If we're not in the middle of progressing through a400 specific constraint `c_i`, we return:401 402 `[t_k1 for k in indices of unfulfilled constraints]`403 404 If we are in the middle of a constraint, then we return:405 `[t_ij]`, where `i` is the index of the inprogress constraint, `j` is the next step for the constraint.406 407 Though we don't care which constraint is fulfilled first, if we are in the progress of fulfilling a constraint,408 that's the only one we'll return.409 """410 token_list = []411 if self.inprogress_constraint is None:412 for constraint in self.pending_constraints: # "pending" == "unfulfilled yet"413 advance = constraint.advance()414 if isinstance(advance, int):415 token_list.append(advance)416 elif isinstance(advance, list):417 token_list.extend(advance)418 else:419 advance = self.inprogress_constraint.advance()420 if isinstance(advance, int):421 token_list.append(advance)422 elif isinstance(advance, list):423 token_list.extend(advance)424 425 if len(token_list) == 0:426 return None427 else:428 return token_list429 430 def reset(self, token_ids: Optional[list[int]]):431 """432 token_ids: the tokens generated thus far to reset the state of the progress through constraints.433 """434 self.init_state()435 436 if token_ids is not None:437 for token in token_ids:438 # completes or steps **one** constraint439 complete, stepped = self.add(token)440 441 # the entire list of constraints are fulfilled442 if self.completed:443 break444 445 def add(self, token_id: int):446 if not isinstance(token_id, int):447 raise TypeError(f"`token_id` should be an `int`, but is `{token_id}`.")448 449 complete, stepped = False, False450 451 if self.completed:452 complete = True453 stepped = False454 return complete, stepped455 456 if self.inprogress_constraint is not None:457 # In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current458 # job, simply update the state459 460 stepped, complete, reset = self.inprogress_constraint.update(token_id)461 if reset:462 # 1. If the next token breaks the progress, then we must restart.463 # e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books".464 465 # But that doesn't mean we self.init_state(), since we only reset the state for this particular466 # constraint, not the full list of constraints.467 468 self.pending_constraints.append(self.inprogress_constraint.copy(stateful=False))469 self.inprogress_constraint = None470 471 if complete:472 # 2. If the next token completes the constraint, move it to completed list, set473 # inprogress to None. If there are no pending constraints either, then this full list of constraints474 # is complete.475 476 self.complete_constraints.append(self.inprogress_constraint)477 self.inprogress_constraint = None478 479 if len(self.pending_constraints) == 0:480 # we're done!481 self.completed = True482 483 else:484 # Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list485 # of constraints?486 487 for cidx, pending_constraint in enumerate(self.pending_constraints):488 if pending_constraint.does_advance(token_id):489 stepped, complete, reset = pending_constraint.update(token_id)490 491 if not stepped:492 raise Exception(493 "`constraint.update(token_id)` is not yielding incremental progress, "494 "even though `constraint.does_advance(token_id)` is true."495 )496 497 if complete:498 self.complete_constraints.append(pending_constraint)499 self.inprogress_constraint = None500 501 if not complete and stepped:502 self.inprogress_constraint = pending_constraint503 504 if complete or stepped:505 # If we made any progress at all, then it's at least not a "pending constraint".506 507 self.pending_constraints = (508 self.pending_constraints[:cidx] + self.pending_constraints[cidx + 1 :]509 )510 511 if len(self.pending_constraints) == 0 and self.inprogress_constraint is None:512 # If there's no longer any pending after this and no inprogress either, then we must be513 # complete.514 515 self.completed = True516 517 break # prevent accidentally stepping through multiple constraints with just one token.518 519 return complete, stepped520 521 def copy(self, stateful=True):522 new_state = ConstraintListState(self.constraints) # we actually never though self.constraints objects523 # throughout this process. So it's at initialization state.524 525 if stateful:526 new_state.complete_constraints = [527 constraint.copy(stateful=True) for constraint in self.complete_constraints528 ]529 if self.inprogress_constraint is not None:530 new_state.inprogress_constraint = self.inprogress_constraint.copy(stateful=True)531 new_state.pending_constraints = [constraint.copy() for constraint in self.pending_constraints]532 533 return new_state534