chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,506 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
"""Implements tracking of constraints for a beam item.
|
||||
|
||||
A list of constraints is given as a list of one or more token
|
||||
sequences, each of length at least one token. For example, for an input sentence
|
||||
|
||||
> Die maschinelle Übersetzung ist schwer zu kontrollieren.
|
||||
|
||||
We could have the constraints:
|
||||
* to influence
|
||||
* hard
|
||||
|
||||
There are two implementations:
|
||||
* OrderedConstraintState: Tracks progress through an ordered list of multitoken constraints.
|
||||
* UnorderedConstraintState: Tracks progress through an unordered list of multitoken constraints.
|
||||
|
||||
The difference is that in the first, the constraints are assumed to be
|
||||
in order; the algorithm will permit zero or more tokens between them.
|
||||
In the second, the constraints are not ordered, so many orderings will
|
||||
be explored.
|
||||
|
||||
The same sequence can be present any number of times, and will appear
|
||||
that many times in the output.
|
||||
"""
|
||||
|
||||
from collections import Counter
|
||||
from typing import List, Optional, Set, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class ConstraintState:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
|
||||
def pack_constraints(batch_constraints: List[List[torch.Tensor]]) -> torch.Tensor:
|
||||
"""Takes a list of list of constraints in tensor form (a list of
|
||||
tensor constraints for each sentence) and transforms it into a
|
||||
packed Tensor. For example, here is a batch of size 3 with 3, 0,
|
||||
and 1 constraints:
|
||||
|
||||
[ [ [3 1 2], [3], [4 5 6 7], ]
|
||||
[],
|
||||
[ [1 8 9 10 1 4 11 12], ]
|
||||
]
|
||||
|
||||
Its corresponding packed structure is:
|
||||
|
||||
[ [ 3 3 1 2 0 3 0 4 5 6 7 0],
|
||||
[ 0 0 0 0 0 0 0 0 0 0 0 0],
|
||||
[ 1 1 8 9 10 1 4 11 12 0 0 0] ]
|
||||
|
||||
The packed tensor has shape (batch size, maxlen), where
|
||||
maxlen is defined below. Each row contains concatenated
|
||||
constraint tokens for that sentence, with 0 appended after
|
||||
each constraint. The first item in each row is the number
|
||||
of constraints for that sentence. So maxlen is the maximum
|
||||
of
|
||||
|
||||
(number of constraints) + (sum length of constraints) + 1.
|
||||
|
||||
across all sentences in the batch.
|
||||
"""
|
||||
# The maximum word length of concatenated constraints for any sentence
|
||||
max_constraints_len = 1
|
||||
for sentence_constraints in batch_constraints:
|
||||
if len(sentence_constraints):
|
||||
# number of constraints, plus sum of constrain lens, plus a zero after each
|
||||
constraints_len = (
|
||||
1
|
||||
+ sum([c.size(0) for c in sentence_constraints])
|
||||
+ len(sentence_constraints)
|
||||
)
|
||||
max_constraints_len = max(max_constraints_len, constraints_len)
|
||||
|
||||
batch_size = len(batch_constraints)
|
||||
constraints_tensor = torch.zeros((batch_size, max_constraints_len)).long()
|
||||
for i, sentence_constraints in enumerate(batch_constraints):
|
||||
constraints_tensor[i, 0] = len(sentence_constraints)
|
||||
offset = 1
|
||||
for j, constraint in enumerate(sentence_constraints):
|
||||
this_len = constraint.size(0)
|
||||
constraints_tensor[i, offset : offset + this_len] = constraint
|
||||
offset += this_len + 1
|
||||
|
||||
return constraints_tensor.long()
|
||||
|
||||
|
||||
def unpack_constraints(constraint_tensor: torch.Tensor) -> List[torch.Tensor]:
|
||||
"""
|
||||
Transforms *one row* of a packed constraint tensor (e.g., for one
|
||||
sentence in the batch) into a list of constraint tensors.
|
||||
"""
|
||||
constraint_list = []
|
||||
num_constraints = constraint_tensor[0]
|
||||
constraints = constraint_tensor.tolist()
|
||||
offset = 1
|
||||
for i in range(num_constraints):
|
||||
where = constraints.index(0, offset)
|
||||
constraint_list.append(constraint_tensor[offset:where])
|
||||
offset = where + 1
|
||||
|
||||
return constraint_list
|
||||
|
||||
|
||||
class ConstraintNode:
|
||||
"""
|
||||
Represents a node in a trie managing unordered constraints.
|
||||
"""
|
||||
|
||||
def __init__(self, token: int = None, parent=None):
|
||||
# The token associate with this node (None for the root)
|
||||
self.token = int(token) if token is not None else None
|
||||
# The parent (None at the root)
|
||||
self.parent = parent
|
||||
# Whether this node is a completed constraint
|
||||
self.terminal = 0
|
||||
# List of child nodes
|
||||
self.children = {}
|
||||
|
||||
# The cumulative number of constraints from this point in the
|
||||
# trie forward
|
||||
self.num_constraints = 0
|
||||
|
||||
@property
|
||||
def id(self):
|
||||
return self.token
|
||||
|
||||
def __str__(self):
|
||||
term = self.terminal != 0
|
||||
return f"[{self.token}].{term}#{self.num_constraints}"
|
||||
|
||||
def __getitem__(self, key: int):
|
||||
return self.children.get(key, None)
|
||||
|
||||
def next_tokens(self) -> Set[int]:
|
||||
"""The set of child labels."""
|
||||
return set(self.children.keys())
|
||||
|
||||
@staticmethod
|
||||
def create(constraints: List[List[int]]):
|
||||
root = ConstraintNode()
|
||||
for sequence in constraints:
|
||||
root.add_sequence(sequence)
|
||||
|
||||
return root
|
||||
|
||||
@staticmethod
|
||||
def print_graph(node: "ConstraintNode"):
|
||||
if len(node.children) == 0:
|
||||
return str(node)
|
||||
else:
|
||||
s = f"({node}"
|
||||
for child in node.children.values():
|
||||
s += " " + ConstraintNode.print_graph(child)
|
||||
s += ")"
|
||||
return s
|
||||
|
||||
def token_counts(self) -> Counter:
|
||||
"""Returns a counter of the number of times each token is used
|
||||
in a constraint.
|
||||
"""
|
||||
token_counts = Counter()
|
||||
kids = list(self.children.values())
|
||||
while len(kids) > 0:
|
||||
kid = kids.pop()
|
||||
token_counts[kid.id] += kid.num_constraints
|
||||
kids += list(kid.children.values())
|
||||
|
||||
return token_counts
|
||||
|
||||
def tokens(self) -> Set[int]:
|
||||
"""Returns the set of tokens in constraints."""
|
||||
return set(self.token_counts().keys())
|
||||
|
||||
def add_sequence(self, sequence: List[int]):
|
||||
"""Adds a constraint, represented as a list of integers, to
|
||||
the trie."""
|
||||
assert len(sequence) > 0
|
||||
|
||||
token = int(sequence[0])
|
||||
if token not in self.children:
|
||||
self.children[token] = ConstraintNode(token, parent=self)
|
||||
|
||||
node = self.children[token]
|
||||
if len(sequence) == 1:
|
||||
node.terminal += 1
|
||||
node.num_constraints += 1
|
||||
parent = node.parent
|
||||
while parent is not None:
|
||||
parent.num_constraints += 1
|
||||
parent = parent.parent
|
||||
else:
|
||||
node.add_sequence(sequence[1:])
|
||||
|
||||
|
||||
class UnorderedConstraintState(ConstraintState):
|
||||
"""
|
||||
Records progress through the set of constraints for each item in the beam
|
||||
using a trie.
|
||||
"""
|
||||
|
||||
def __init__(self, node: ConstraintNode, copy_from: "ConstraintState" = None):
|
||||
self.node = node
|
||||
|
||||
if copy_from is None:
|
||||
# The root node
|
||||
self.root = node
|
||||
# The set of states in the graph that have been completed
|
||||
self.completed = Counter()
|
||||
# The...
|
||||
self.generated = Counter()
|
||||
# The list of tokens we need to generate
|
||||
self.needed_tokens = self.root.tokens()
|
||||
else:
|
||||
self.completed = Counter(copy_from.completed)
|
||||
self.generated = Counter(copy_from.generated)
|
||||
self.root = copy_from.root
|
||||
|
||||
# Mark the node as generated
|
||||
if self.node != self.root:
|
||||
self.generated[node] += 1
|
||||
|
||||
@staticmethod
|
||||
def create(constraint_tensor: torch.Tensor):
|
||||
constraint_list = unpack_constraints(constraint_tensor)
|
||||
constraint_trie_root = ConstraintNode.create(constraint_list)
|
||||
return UnorderedConstraintState(constraint_trie_root)
|
||||
|
||||
def __str__(self):
|
||||
gen_str = ",".join([str(node) for node in self.generated])
|
||||
return f"{self.name}/{self.bank}({gen_str})x{self.num_completed}"
|
||||
|
||||
def __copy__(self):
|
||||
copied_state = UnorderedConstraintState(self.node, copy_from=self)
|
||||
return copied_state
|
||||
|
||||
def copy(self):
|
||||
return self.__copy__()
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
if self.node.id is None:
|
||||
return "ROOT"
|
||||
else:
|
||||
return str(self.node.id)
|
||||
|
||||
@property
|
||||
def is_root(self):
|
||||
return self.node == self.root
|
||||
|
||||
@property
|
||||
def bank(self):
|
||||
return sum(self.generated.values())
|
||||
|
||||
@property
|
||||
def num_completed(self):
|
||||
"""The number of constraints (not constraint tokens) that are completed.
|
||||
In addition to the already-completed states, we need to account for the
|
||||
current state, which might get marked as completed when another token
|
||||
is generated.
|
||||
"""
|
||||
in_final = self.node.terminal and self.completed[self.node] < self.node.terminal
|
||||
return sum(self.completed.values()) + in_final
|
||||
|
||||
@property
|
||||
def finished(self):
|
||||
return self.root.num_constraints - self.num_completed == 0
|
||||
|
||||
@property
|
||||
def token_counts(self):
|
||||
return self.root.token_counts()
|
||||
|
||||
@property
|
||||
def tokens(self):
|
||||
return self.root.tokens()
|
||||
|
||||
@property
|
||||
def num_constraint_tokens(self):
|
||||
return sum(self.token_counts.values())
|
||||
|
||||
def next_tokens(self) -> Set[int]:
|
||||
"""Returns the list of tokens that could come next.
|
||||
These are (a) all tokens extending the root state and, for
|
||||
non-root states, additionally all tokens extending the current
|
||||
state."""
|
||||
|
||||
if self.node != self.root:
|
||||
return self.root.next_tokens().union(self.node.next_tokens())
|
||||
else:
|
||||
return self.root.next_tokens()
|
||||
|
||||
def advance(self, token: int):
|
||||
"""Reads in a token and advances the state. Here's how it works.
|
||||
|
||||
We can advance to the next state if:
|
||||
- there is a matching child
|
||||
- its path isn't blocked
|
||||
|
||||
A path is blocked when all constraints that are descendants of
|
||||
that node have already been generated, in the current state.
|
||||
|
||||
If we are not able to advance from the current state, we "fall
|
||||
off the graph" and return to the root state. There, we again
|
||||
try to advance, checking the same criteria.
|
||||
|
||||
In any case, when falling off the graph, we need to do some
|
||||
bookkeeping. We:
|
||||
- check whether any constraints were met (all prefixes of
|
||||
current state)
|
||||
- if one is found, mark it as completed
|
||||
- adjust visited nodes accordingly
|
||||
"""
|
||||
token = int(token)
|
||||
|
||||
next_state = None
|
||||
child = self.node[token]
|
||||
if child is not None and self.generated[child] < child.num_constraints:
|
||||
next_state = UnorderedConstraintState(child, copy_from=self)
|
||||
|
||||
def rewind():
|
||||
"""If we're mid-trie and an "illegal" token is chosen next, we need
|
||||
to reset our state to the root state. However, along the way, we need
|
||||
to check whether a prefix of the current trie state represents a state
|
||||
we could mark as completed.
|
||||
"""
|
||||
node = self.node
|
||||
while node != self.root:
|
||||
if node.terminal and self.completed[node] < node.terminal:
|
||||
next_state.completed[node] += 1
|
||||
return
|
||||
|
||||
next_state.generated[node] -= 1
|
||||
node = node.parent
|
||||
|
||||
# Fall off the graph, check the root
|
||||
if next_state is None and token in self.root.next_tokens():
|
||||
child = self.root[token]
|
||||
# We can only traverse this edge if it's not saturated
|
||||
if self.generated[child] < child.num_constraints:
|
||||
next_state = UnorderedConstraintState(child, copy_from=self)
|
||||
else:
|
||||
next_state = UnorderedConstraintState(self.root, copy_from=self)
|
||||
|
||||
# Rewind
|
||||
rewind()
|
||||
|
||||
elif next_state is None:
|
||||
next_state = UnorderedConstraintState(self.root, copy_from=self)
|
||||
# Rewind
|
||||
rewind()
|
||||
|
||||
return next_state
|
||||
|
||||
|
||||
class ConstraintSequence:
|
||||
def __init__(self, sequences: List[List[int]]):
|
||||
"""Represents a set of possibly multitoken constraints by
|
||||
concatenating them and internally recording the end points.
|
||||
"""
|
||||
self.sequences = []
|
||||
self.endpoints = []
|
||||
self.num_tokens = 0
|
||||
self.tokens = set()
|
||||
for sequence in sequences:
|
||||
for token in sequence:
|
||||
self.tokens.add(token)
|
||||
self.num_tokens += len(sequence)
|
||||
self.endpoints += [False for x in range(len(sequence) - 1)] + [True]
|
||||
self.sequences += sequence
|
||||
|
||||
def __getitem__(self, key: int):
|
||||
return self.sequences[key]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.sequences)
|
||||
|
||||
def __str__(self):
|
||||
return str(self.sequences)
|
||||
|
||||
|
||||
class OrderedConstraintState(ConstraintState):
|
||||
"""
|
||||
Records progress through the set of linear nonbranching constraints with gaps.
|
||||
"""
|
||||
|
||||
def __init__(self, sequence: ConstraintSequence, state: int = -1):
|
||||
self.sequence = sequence
|
||||
self.state = state
|
||||
|
||||
@staticmethod
|
||||
def create(constraint_tensor: torch.Tensor):
|
||||
constraint_list = unpack_constraints(constraint_tensor)
|
||||
return OrderedConstraintState(ConstraintSequence(constraint_list), -1)
|
||||
|
||||
def __str__(self):
|
||||
return f"{self.state}/{self.bank}x{self.num_completed}"
|
||||
|
||||
def __copy__(self):
|
||||
return OrderedConstraintState(self.sequence, self.state)
|
||||
|
||||
def copy(self):
|
||||
return self.__copy__()
|
||||
|
||||
@property
|
||||
def num_completed(self):
|
||||
if self.state == -1:
|
||||
return 0
|
||||
count = len(
|
||||
list(filter(lambda x: x, self.sequence.endpoints[0 : self.state + 1]))
|
||||
)
|
||||
return count
|
||||
|
||||
@property
|
||||
def is_root(self):
|
||||
return self.state == -1
|
||||
|
||||
@property
|
||||
def name(self):
|
||||
if self.state == -1:
|
||||
return "ROOT"
|
||||
else:
|
||||
return str(self.sequence[self.state])
|
||||
|
||||
@property
|
||||
def bank(self) -> int:
|
||||
return self.state + 1
|
||||
|
||||
@property
|
||||
def finished(self):
|
||||
return self.state + 1 == len(self.sequence)
|
||||
|
||||
@property
|
||||
def token_counts(self):
|
||||
return self.sequence.token_counts()
|
||||
|
||||
@property
|
||||
def tokens(self):
|
||||
return self.sequence.tokens
|
||||
|
||||
@property
|
||||
def num_constraint_tokens(self):
|
||||
return sum(self.token_counts.values())
|
||||
|
||||
def next_tokens(self) -> Set[int]:
|
||||
"""Returns the list of tokens that could come next.
|
||||
These are (a) all tokens extending the root state and, for
|
||||
non-root states, additionally all tokens extending the current
|
||||
state."""
|
||||
|
||||
tokens = set()
|
||||
if self.state > 0:
|
||||
tokens.add(self.sequence[0])
|
||||
if not self.finished:
|
||||
tokens.add(self.sequence[self.state + 1])
|
||||
return tokens
|
||||
|
||||
def advance(self, token: int):
|
||||
"""Reads in a token and advances the state. Here's how it works.
|
||||
|
||||
We can advance to the next state if:
|
||||
- there is a matching child
|
||||
- its path isn't blocked
|
||||
|
||||
A path is blocked when all constraints that are descendants of
|
||||
that node have already been generated, in the current state.
|
||||
|
||||
If we are not able to advance from the current state, we "fall
|
||||
off the graph" and return to the root state. There, we again
|
||||
try to advance, checking the same criteria.
|
||||
|
||||
In any case, when falling off the graph, we need to do some
|
||||
bookkeeping. We:
|
||||
- check whether any constraints were met (all prefixes of
|
||||
current state)
|
||||
- if one is found, mark it as completed
|
||||
- adjust visited nodes accordingly
|
||||
"""
|
||||
token = int(token)
|
||||
# print(f"{self} ADVANCE({token}) {self.sequence} -> ", end="")
|
||||
|
||||
if self.finished:
|
||||
# Accept anything
|
||||
next_state = self.copy()
|
||||
|
||||
elif self.sequence[self.state + 1] == token:
|
||||
# Advance to the next token
|
||||
next_state = OrderedConstraintState(self.sequence, self.state + 1)
|
||||
|
||||
elif self.sequence.endpoints[self.state]:
|
||||
# Accept anything between constraints (*)
|
||||
next_state = self.copy()
|
||||
|
||||
elif token == self.sequence[0]:
|
||||
# Start over having generated the first token
|
||||
next_state = OrderedConstraintState(self.sequence, 0)
|
||||
else:
|
||||
# Start over from the root
|
||||
next_state = OrderedConstraintState(self.sequence, -1)
|
||||
|
||||
return next_state
|
||||
Reference in New Issue
Block a user