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