def __init__(self, logits, l, bound):
self.bound = bound
self.logits = logits
argmax = tf.argmax(logits, axis=1)
self.best_score = tf.reduce_max(logits, axis=1)
self.predicted_span = to_unpacked_coordinates(argmax, l, bound)
self.l = l
def get_best_span(self, bound):
if bound > self.bound:
raise ValueError()
if bound < self.bound:
cutoff = self.l * bound - bound * (bound - 1) // 2
logits = self.logits[:, :cutoff]
argmax = tf.argmax(logits, axis=1)
best_score = tf.reduce_max(logits, axis=1)
predicted_span = to_unpacked_coordinates(argmax, self.l, bound)
return predicted_span, best_score
return self.predicted_span, self.best_score