跳转至

beam_constraints

mindnlp.transformers.generation.beam_constraints

Beam constraints

mindnlp.transformers.generation.beam_constraints.Constraint

Bases: ABC

Abstract base class for all constraints that can be applied during generation. It must define how the constraint can be satisfied.

All classes that inherit Constraint must follow the requirement that

completed = False
while not completed:
    _, completed = constraint.update(constraint.advance())

will always terminate (halt).

Source code in mindnlp\transformers\generation\beam_constraints.py 
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
class Constraint(ABC):
    r"""Abstract base class for all constraints that can be applied during generation.
    It must define how the constraint can be satisfied.

    All classes that inherit Constraint must follow the requirement that

    ```py
    completed = False
    while not completed:
        _, completed = constraint.update(constraint.advance())
    ```

    will always terminate (halt).
    """

    def __init__(self):
        # test for the above condition
        self.test()

    def test(self):
        """
        Tests whether this constraint has been properly defined.
        """
        counter = 0
        completed = False
        while not completed:
            if counter == 1:
                self.reset()
            advance = self.advance()
            if not self.does_advance(advance):
                raise Exception(
                    "Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true."
                )

            stepped, completed, reset = self.update(advance)
            counter += 1

            if counter > 10000:
                raise Exception("update() does not fulfill the constraint.")

        if self.remaining() != 0:
            raise Exception("Custom Constraint is not defined correctly.")

    @abstractmethod
    def advance(self):
        """
        When called, returns the token that would take this constraint one step closer to being fulfilled.

        Return:
            token_ids(`torch.tensor`): Must be a tensor of a list of indexable tokens, not some integer.
        """
        raise NotImplementedError(
            f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
        )

    @abstractmethod
    def does_advance(self, token_id: int):
        """
        Reads in a token and returns whether it creates progress.
        """
        raise NotImplementedError(
            f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
        )

    @abstractmethod
    def update(self, token_id: int):
        """
        Reads in a token and returns booleans that indicate the progress made by it. This function will update the
        state of this object unlikes `does_advance(self, token_id: int)`.

        This isn't to test whether a certain token will advance the progress; it's to update its state as if it has
        been generated. This becomes important if token_id != desired token (refer to else statement in
        PhrasalConstraint)

        Args:
            token_id(`int`):
                The id of a newly generated token in the beam search.
        Return:
            stepped(`bool`):
                Whether this constraint has become one step closer to being fulfuilled.
            completed(`bool`):
                Whether this constraint has been completely fulfilled by this token being generated.
            reset (`bool`):
                Whether this constraint has reset its progress by this token being generated.
        """
        raise NotImplementedError(
            f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
        )

    @abstractmethod
    def reset(self):
        """
        Resets the state of this constraint to its initialization. We would call this in cases where the fulfillment of
        a constraint is abrupted by an unwanted token.
        """
        raise NotImplementedError(
            f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
        )

    @abstractmethod
    def remaining(self):
        """
        Returns the number of remaining steps of `advance()` in order to complete this constraint.
        """
        raise NotImplementedError(
            f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
        )

    @abstractmethod
    def copy(self, stateful=False):
        """
        Creates a new instance of this constraint.

        Args:
            stateful(`bool`): Whether to not only copy the constraint for new instance, but also its state.

        Return:
            constraint(`Constraint`): The same constraint as the one being called from.
        """
        raise NotImplementedError(
            f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
        )

mindnlp.transformers.generation.beam_constraints.Constraint.advance() abstractmethod

When called, returns the token that would take this constraint one step closer to being fulfilled.

Return

token_ids(torch.tensor): Must be a tensor of a list of indexable tokens, not some integer.

Source code in mindnlp\transformers\generation\beam_constraints.py 
65
66
67
68
69
70
71
72
73
74
75
@abstractmethod
def advance(self):
    """
    When called, returns the token that would take this constraint one step closer to being fulfilled.

    Return:
        token_ids(`torch.tensor`): Must be a tensor of a list of indexable tokens, not some integer.
    """
    raise NotImplementedError(
        f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
    )

mindnlp.transformers.generation.beam_constraints.Constraint.copy(stateful=False) abstractmethod

Creates a new instance of this constraint.

PARAMETER DESCRIPTION
stateful(`bool`)

Whether to not only copy the constraint for new instance, but also its state.
Return

constraint(Constraint): The same constraint as the one being called from.

Source code in mindnlp\transformers\generation\beam_constraints.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
@abstractmethod
def copy(self, stateful=False):
    """
    Creates a new instance of this constraint.

    Args:
        stateful(`bool`): Whether to not only copy the constraint for new instance, but also its state.

    Return:
        constraint(`Constraint`): The same constraint as the one being called from.
    """
    raise NotImplementedError(
        f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
    )

mindnlp.transformers.generation.beam_constraints.Constraint.does_advance(token_id) abstractmethod

Reads in a token and returns whether it creates progress.

Source code in mindnlp\transformers\generation\beam_constraints.py 
77
78
79
80
81
82
83
84
@abstractmethod
def does_advance(self, token_id: int):
    """
    Reads in a token and returns whether it creates progress.
    """
    raise NotImplementedError(
        f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
    )

mindnlp.transformers.generation.beam_constraints.Constraint.remaining() abstractmethod

Returns the number of remaining steps of advance() in order to complete this constraint.

Source code in mindnlp\transformers\generation\beam_constraints.py 
121
122
123
124
125
126
127
128
@abstractmethod
def remaining(self):
    """
    Returns the number of remaining steps of `advance()` in order to complete this constraint.
    """
    raise NotImplementedError(
        f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
    )

mindnlp.transformers.generation.beam_constraints.Constraint.reset() abstractmethod

Resets the state of this constraint to its initialization. We would call this in cases where the fulfillment of a constraint is abrupted by an unwanted token.

Source code in mindnlp\transformers\generation\beam_constraints.py 
111
112
113
114
115
116
117
118
119
@abstractmethod
def reset(self):
    """
    Resets the state of this constraint to its initialization. We would call this in cases where the fulfillment of
    a constraint is abrupted by an unwanted token.
    """
    raise NotImplementedError(
        f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
    )

mindnlp.transformers.generation.beam_constraints.Constraint.test()

Tests whether this constraint has been properly defined.

Source code in mindnlp\transformers\generation\beam_constraints.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def test(self):
    """
    Tests whether this constraint has been properly defined.
    """
    counter = 0
    completed = False
    while not completed:
        if counter == 1:
            self.reset()
        advance = self.advance()
        if not self.does_advance(advance):
            raise Exception(
                "Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true."
            )

        stepped, completed, reset = self.update(advance)
        counter += 1

        if counter > 10000:
            raise Exception("update() does not fulfill the constraint.")

    if self.remaining() != 0:
        raise Exception("Custom Constraint is not defined correctly.")

mindnlp.transformers.generation.beam_constraints.Constraint.update(token_id) abstractmethod

Reads in a token and returns booleans that indicate the progress made by it. This function will update the state of this object unlikes does_advance(self, token_id: int).

This isn't to test whether a certain token will advance the progress; it's to update its state as if it has been generated. This becomes important if token_id != desired token (refer to else statement in PhrasalConstraint)

PARAMETER DESCRIPTION
token_id(`int`)

The id of a newly generated token in the beam search.
Source code in mindnlp\transformers\generation\beam_constraints.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
@abstractmethod
def update(self, token_id: int):
    """
    Reads in a token and returns booleans that indicate the progress made by it. This function will update the
    state of this object unlikes `does_advance(self, token_id: int)`.

    This isn't to test whether a certain token will advance the progress; it's to update its state as if it has
    been generated. This becomes important if token_id != desired token (refer to else statement in
    PhrasalConstraint)

    Args:
        token_id(`int`):
            The id of a newly generated token in the beam search.
    Return:
        stepped(`bool`):
            Whether this constraint has become one step closer to being fulfuilled.
        completed(`bool`):
            Whether this constraint has been completely fulfilled by this token being generated.
        reset (`bool`):
            Whether this constraint has reset its progress by this token being generated.
    """
    raise NotImplementedError(
        f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
    )

mindnlp.transformers.generation.beam_constraints.ConstraintListState

A class for beam scorers to track its progress through a list of constraints.

PARAMETER DESCRIPTION
constraints

A list of [Constraint] objects that must be fulfilled by the beam scorer.

TYPE: `List[Constraint]`

Source code in mindnlp\transformers\generation\beam_constraints.py 
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
class ConstraintListState:
    r"""
    A class for beam scorers to track its progress through a list of constraints.

    Args:
        constraints (`List[Constraint]`):
            A list of [`Constraint`] objects that must be fulfilled by the beam scorer.
    """

    def __init__(self, constraints: List[Constraint]):
        self.constraints = constraints

        # max # of steps required to fulfill a given constraint
        self.max_seqlen = max(c.seqlen for c in constraints)
        self.n_constraints = len(constraints)
        self.completed = False

        self.init_state()

    def init_state(self):
        self.complete_constraints = []
        self.inprogress_constraint = None
        self.pending_constraints = [constraint.copy(stateful=False) for constraint in self.constraints]

    def get_bank(self):
        add = 0
        if self.inprogress_constraint:
            # extra points for having a constraint mid-fulfilled
            add += self.max_seqlen - self.inprogress_constraint.remaining()

        return (len(self.complete_constraints) * self.max_seqlen) + add

    def advance(self):
        """The list of tokens to generate such that we can make progress.
        By "list" we don't mean the list of token that will fully fulfill a constraint.

        Given constraints `c_i = {t_ij | j == # of tokens}`, If we're not in the middle of progressing through a
        specific constraint `c_i`, we return:

        `[t_k1 for k in indices of unfulfilled constraints]`

        If we are in the middle of a constraint, then we return:
            `[t_ij]`, where `i` is the index of the inprogress constraint, `j` is the next step for the constraint.

        Though we don't care which constraint is fulfilled first, if we are in the progress of fulfilling a constraint,
        that's the only one we'll return.
        """
        token_list = []
        if self.inprogress_constraint is None:
            for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
                advance = constraint.advance()
                if isinstance(advance, int):
                    token_list.append(advance)
                elif isinstance(advance, list):
                    token_list.extend(advance)
        else:
            advance = self.inprogress_constraint.advance()
            if isinstance(advance, int):
                token_list.append(advance)
            elif isinstance(advance, list):
                token_list.extend(advance)

        if len(token_list) == 0:
            return None
        else:
            return token_list

    def reset(self, token_ids: Optional[List[int]]):
        """
        token_ids: the tokens generated thus far to reset the state of the progress through constraints.
        """
        self.init_state()

        if token_ids is not None:
            for token in token_ids:
                # completes or steps **one** constraint
                complete, stepped = self.add(token)

                # the entire list of constraints are fulfilled
                if self.completed:
                    break

    def add(self, token_id: int):
        if not isinstance(token_id, int):
            raise TypeError(f"`token_id` should be an `int`, but is `{token_id}`.")

        complete, stepped = False, False

        if self.completed:
            complete = True
            stepped = False
            return complete, stepped

        if self.inprogress_constraint is not None:
            # In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current
            # job, simply update the state

            stepped, complete, reset = self.inprogress_constraint.update(token_id)
            if reset:
                # 1. If the next token breaks the progress, then we must restart.
                #     e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books".

                #     But that doesn't mean we self.init_state(), since we only reset the state for this particular
                #     constraint, not the full list of constraints.

                self.pending_constraints.append(self.inprogress_constraint.copy(stateful=False))
                self.inprogress_constraint = None

            if complete:
                # 2. If the next token completes the constraint, move it to completed list, set
                #     inprogress to None. If there are no pending constraints either, then this full list of constraints
                #     is complete.

                self.complete_constraints.append(self.inprogress_constraint)
                self.inprogress_constraint = None

                if len(self.pending_constraints) == 0:
                    # we're done!
                    self.completed = True

        else:
            # Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list
            # of constraints?

            for cidx, pending_constraint in enumerate(self.pending_constraints):
                if pending_constraint.does_advance(token_id):
                    stepped, complete, reset = pending_constraint.update(token_id)

                    if not stepped:
                        raise Exception(
                            "`constraint.update(token_id)` is not yielding incremental progress, "
                            "even though `constraint.does_advance(token_id)` is true."
                        )

                    if complete:
                        self.complete_constraints.append(pending_constraint)
                        self.inprogress_constraint = None

                    if not complete and stepped:
                        self.inprogress_constraint = pending_constraint

                    if complete or stepped:
                        # If we made any progress at all, then it's at least not a "pending constraint".

                        self.pending_constraints = (
                            self.pending_constraints[:cidx] + self.pending_constraints[cidx + 1 :]
                        )

                        if len(self.pending_constraints) == 0 and self.inprogress_constraint is None:
                            # If there's no longer any pending after this and no inprogress either, then we must be
                            # complete.

                            self.completed = True

                        break  # prevent accidentally stepping through multiple constraints with just one token.

        return complete, stepped

    def copy(self, stateful=True):
        new_state = ConstraintListState(self.constraints)  # we actually never though self.constraints objects
        # throughout this process. So it's at initialization state.

        if stateful:
            new_state.complete_constraints = [
                constraint.copy(stateful=True) for constraint in self.complete_constraints
            ]
            if self.inprogress_constraint is not None:
                new_state.inprogress_constraint = self.inprogress_constraint.copy(stateful=True)
            new_state.pending_constraints = [constraint.copy() for constraint in self.pending_constraints]

        return new_state

mindnlp.transformers.generation.beam_constraints.ConstraintListState.advance()

The list of tokens to generate such that we can make progress. By "list" we don't mean the list of token that will fully fulfill a constraint.

Given constraints c_i = {t_ij | j == # of tokens}, If we're not in the middle of progressing through a specific constraint c_i, we return:

[t_k1 for k in indices of unfulfilled constraints]

If we are in the middle of a constraint, then we return: [t_ij], where i is the index of the inprogress constraint, j is the next step for the constraint.

Though we don't care which constraint is fulfilled first, if we are in the progress of fulfilling a constraint, that's the only one we'll return.

Source code in mindnlp\transformers\generation\beam_constraints.py 
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
def advance(self):
    """The list of tokens to generate such that we can make progress.
    By "list" we don't mean the list of token that will fully fulfill a constraint.

    Given constraints `c_i = {t_ij | j == # of tokens}`, If we're not in the middle of progressing through a
    specific constraint `c_i`, we return:

    `[t_k1 for k in indices of unfulfilled constraints]`

    If we are in the middle of a constraint, then we return:
        `[t_ij]`, where `i` is the index of the inprogress constraint, `j` is the next step for the constraint.

    Though we don't care which constraint is fulfilled first, if we are in the progress of fulfilling a constraint,
    that's the only one we'll return.
    """
    token_list = []
    if self.inprogress_constraint is None:
        for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
            advance = constraint.advance()
            if isinstance(advance, int):
                token_list.append(advance)
            elif isinstance(advance, list):
                token_list.extend(advance)
    else:
        advance = self.inprogress_constraint.advance()
        if isinstance(advance, int):
            token_list.append(advance)
        elif isinstance(advance, list):
            token_list.extend(advance)

    if len(token_list) == 0:
        return None
    else:
        return token_list

mindnlp.transformers.generation.beam_constraints.ConstraintListState.reset(token_ids)

Source code in mindnlp\transformers\generation\beam_constraints.py 
435
436
437
438
439
440
441
442
443
444
445
446
447
448
def reset(self, token_ids: Optional[List[int]]):
    """
    token_ids: the tokens generated thus far to reset the state of the progress through constraints.
    """
    self.init_state()

    if token_ids is not None:
        for token in token_ids:
            # completes or steps **one** constraint
            complete, stepped = self.add(token)

            # the entire list of constraints are fulfilled
            if self.completed:
                break

mindnlp.transformers.generation.beam_constraints.DisjunctiveConstraint

Bases: Constraint

A special [Constraint] that is fulfilled by fulfilling just one of several constraints.

PARAMETER DESCRIPTION
nested_token_ids

A list of words, where each word is a list of ids. This constraint is fulfilled by generating just one from the list of words.

TYPE: `List[List[int]]`

Source code in mindnlp\transformers\generation\beam_constraints.py 
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
class DisjunctiveConstraint(Constraint):
    r"""
    A special [`Constraint`] that is fulfilled by fulfilling just one of several constraints.

    Args:
        nested_token_ids (`List[List[int]]`):
            A list of words, where each word is a list of ids. This constraint is fulfilled by generating just one from
            the list of words.
    """

    def __init__(self, nested_token_ids: List[List[int]]):
        super(Constraint, self).__init__()

        if not isinstance(nested_token_ids, list) or len(nested_token_ids) == 0:
            raise ValueError(f"`nested_token_ids` has to be a non-empty list, but is {nested_token_ids}.")
        if any(not isinstance(token_ids, list) for token_ids in nested_token_ids):
            raise ValueError(f"`nested_token_ids` has to be a list of lists, but is {nested_token_ids}.")
        if any(
            any((not isinstance(token_id, int) or token_id < 0) for token_id in token_ids)
            for token_ids in nested_token_ids
        ):
            raise ValueError(
                f"Each list in `nested_token_ids` has to be a list of positive integers, but is {nested_token_ids}."
            )

        self.trie = DisjunctiveTrie(nested_token_ids)
        self.token_ids = nested_token_ids

        self.seqlen = self.trie.max_height
        self.current_seq = []
        self.completed = False

    def advance(self):
        token_list = self.trie.next_tokens(self.current_seq)

        if len(token_list) == 0:
            return None
        else:
            return token_list

    def does_advance(self, token_id: int):
        if not isinstance(token_id, int):
            raise TypeError(f"`token_id` is supposed to be type `int`, but is {token_id} of type {type(token_id)}")

        next_tokens = self.trie.next_tokens(self.current_seq)

        return token_id in next_tokens

    def update(self, token_id: int):
        if not isinstance(token_id, int):
            raise TypeError(f"`token_id` is supposed to be type `int`, but is {token_id} of type {type(token_id)}")

        stepped = False
        completed = False
        reset = False

        if self.does_advance(token_id):
            self.current_seq.append(token_id)
            stepped = True
        else:
            reset = True
            self.reset()

        completed = self.trie.reached_leaf(self.current_seq)
        self.completed = completed

        return stepped, completed, reset

    def reset(self):
        self.completed = False
        self.current_seq = []

    def remaining(self):
        if self.completed:
            # since this can be completed without reaching max height
            return 0
        else:
            return self.seqlen - len(self.current_seq)

    def copy(self, stateful=False):
        new_constraint = DisjunctiveConstraint(self.token_ids)

        if stateful:
            new_constraint.seq_len = self.seqlen
            new_constraint.current_seq = self.current_seq
            new_constraint.completed = self.completed

        return new_constraint

mindnlp.transformers.generation.beam_constraints.DisjunctiveTrie

Source code in mindnlp\transformers\generation\beam_constraints.py 
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
class DisjunctiveTrie:
    def __init__(self, nested_token_ids: List[List[int]], no_subsets=True):
        r"""
        A helper class that builds a trie with the words represented in `nested_token_ids`.
        """
        self.max_height = max(len(one) for one in nested_token_ids)

        root = {}
        for token_ids in nested_token_ids:
            level = root
            for tidx, token_id in enumerate(token_ids):
                if token_id not in level:
                    level[token_id] = {}

                level = level[token_id]

        if no_subsets and self.has_subsets(root, nested_token_ids):
            raise ValueError(
                "Each list in `nested_token_ids` can't be a complete subset of another list, but is"
                f" {nested_token_ids}."
            )

        self.trie = root

    def next_tokens(self, current_seq):
        """
        The next possible tokens that will progress the trie, given the current sequence of tokens in `current_seq`.
        """
        start = self.trie

        for current_token in current_seq:
            start = start[current_token]

        next_tokens = list(start.keys())

        return next_tokens

    def reached_leaf(self, current_seq):
        next_tokens = self.next_tokens(current_seq)

        return len(next_tokens) == 0

    def count_leaves(self, root):
        next_nodes = list(root.values())
        if len(next_nodes) == 0:
            return 1
        else:
            return sum(self.count_leaves(nn) for nn in next_nodes)

    def has_subsets(self, trie, nested_token_ids):
        """
        Returns whether # of leaves == # of words. Otherwise some word is a subset of another.
        """
        leaf_count = self.count_leaves(trie)
        return len(nested_token_ids) != leaf_count

mindnlp.transformers.generation.beam_constraints.DisjunctiveTrie.__init__(nested_token_ids, no_subsets=True)

A helper class that builds a trie with the words represented in nested_token_ids.

Source code in mindnlp\transformers\generation\beam_constraints.py 
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
def __init__(self, nested_token_ids: List[List[int]], no_subsets=True):
    r"""
    A helper class that builds a trie with the words represented in `nested_token_ids`.
    """
    self.max_height = max(len(one) for one in nested_token_ids)

    root = {}
    for token_ids in nested_token_ids:
        level = root
        for tidx, token_id in enumerate(token_ids):
            if token_id not in level:
                level[token_id] = {}

            level = level[token_id]

    if no_subsets and self.has_subsets(root, nested_token_ids):
        raise ValueError(
            "Each list in `nested_token_ids` can't be a complete subset of another list, but is"
            f" {nested_token_ids}."
        )

    self.trie = root

mindnlp.transformers.generation.beam_constraints.DisjunctiveTrie.has_subsets(trie, nested_token_ids)

Returns whether # of leaves == # of words. Otherwise some word is a subset of another.

Source code in mindnlp\transformers\generation\beam_constraints.py 
270
271
272
273
274
275
def has_subsets(self, trie, nested_token_ids):
    """
    Returns whether # of leaves == # of words. Otherwise some word is a subset of another.
    """
    leaf_count = self.count_leaves(trie)
    return len(nested_token_ids) != leaf_count

mindnlp.transformers.generation.beam_constraints.DisjunctiveTrie.next_tokens(current_seq)

The next possible tokens that will progress the trie, given the current sequence of tokens in current_seq.

Source code in mindnlp\transformers\generation\beam_constraints.py 
245
246
247
248
249
250
251
252
253
254
255
256
def next_tokens(self, current_seq):
    """
    The next possible tokens that will progress the trie, given the current sequence of tokens in `current_seq`.
    """
    start = self.trie

    for current_token in current_seq:
        start = start[current_token]

    next_tokens = list(start.keys())

    return next_tokens

mindnlp.transformers.generation.beam_constraints.PhrasalConstraint

Bases: Constraint

[Constraint] enforcing that an ordered sequence of tokens is included in the output.

PARAMETER DESCRIPTION
token_ids

The id of the token that must be generated by the output.

TYPE: `List[int]`

Source code in mindnlp\transformers\generation\beam_constraints.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
class PhrasalConstraint(Constraint):
    r"""
    [`Constraint`] enforcing that an ordered sequence of tokens is included in the output.

    Args:
        token_ids (`List[int]`):
            The id of the token that must be generated by the output.
    """

    def __init__(self, token_ids: List[int]):
        super(Constraint, self).__init__()

        if not isinstance(token_ids, list) or len(token_ids) == 0:
            raise ValueError(f"`token_ids` has to be a non-empty list, but is {token_ids}.")
        if any((not isinstance(token_id, int) or token_id < 0) for token_id in token_ids):
            raise ValueError(f"Each list in `token_ids` has to be a list of positive integers, but is {token_ids}.")

        self.token_ids = token_ids

        self.seqlen = len(self.token_ids)
        self.fulfilled_idx = -1  # the index of the currently fulfilled step
        self.completed = False

    def advance(self):
        if self.completed:
            return None
        return self.token_ids[self.fulfilled_idx + 1]

    def does_advance(self, token_id: int):
        if not isinstance(token_id, int):
            raise TypeError(f"`token_id` has to be an `int`, but is {token_id} of type {type(token_id)}")

        if self.completed:
            return False

        return token_id == self.token_ids[self.fulfilled_idx + 1]

    def update(self, token_id: int):
        if not isinstance(token_id, int):
            raise TypeError(f"`token_id` has to be an `int`, but is {token_id} of type {type(token_id)}")

        stepped = False
        completed = False
        reset = False

        if self.does_advance(token_id):
            self.fulfilled_idx += 1
            stepped = True
            if self.fulfilled_idx == (self.seqlen - 1):
                completed = True
            self.completed = completed
        else:
            # failed to make progress.
            reset = True
            self.reset()
        return stepped, completed, reset

    def reset(self):
        self.completed = False
        self.fulfilled_idx = 0

    def remaining(self):
        return self.seqlen - (self.fulfilled_idx + 1)

    def copy(self, stateful=False):
        new_constraint = PhrasalConstraint(self.token_ids)

        if stateful:
            new_constraint.seq_len = self.seqlen
            new_constraint.fulfilled_idx = self.fulfilled_idx
            new_constraint.completed = self.completed

        return new_constraint