Better approach for handling zero suggestions

This commit is contained in:
Lj Miranda 2022-12-21 20:01:02 +08:00
parent a3fad0b983
commit 8c4eee28bc

View File

@ -266,26 +266,27 @@ class Exclusive_SpanCategorizer(SpanCategorizer):
) -> SpanGroup: ) -> SpanGroup:
scores = self.model.ops.to_numpy(scores) scores = self.model.ops.to_numpy(scores)
indices = self.model.ops.to_numpy(indices) indices = self.model.ops.to_numpy(indices)
if scores.size != 0:
predicted = scores.argmax(axis=1)
# Remove samples where the negative label is the argmax # Handle cases when there are zero suggestions
positive = numpy.where(predicted != self._negative_label)[0] if scores.size == 0:
predicted = predicted[positive] return SpanGroup(doc, name=self.key)
indices = indices[positive]
# Sort spans according to argmax probability predicted = scores.argmax(axis=1)
if not allow_overlap and predicted.size != 0: # Remove samples where the negative label is the argmax
# Get the probabilities positive = numpy.where(predicted != self._negative_label)[0]
argmax_probs = numpy.take_along_axis( predicted = predicted[positive]
scores[positive], numpy.expand_dims(predicted, 1), axis=1 indices = indices[positive]
)
argmax_probs = argmax_probs.squeeze() # Sort spans according to argmax probability
sort_idx = (argmax_probs * -1).argsort() if not allow_overlap and predicted.size != 0:
predicted = predicted[sort_idx] # Get the probabilities
indices = indices[sort_idx] argmax_probs = numpy.take_along_axis(
else: scores[positive], numpy.expand_dims(predicted, 1), axis=1
predicted = [] )
argmax_probs = argmax_probs.squeeze()
sort_idx = (argmax_probs * -1).argsort()
predicted = predicted[sort_idx]
indices = indices[sort_idx]
seen = Ranges() seen = Ranges()
spans = SpanGroup(doc, name=self.key) spans = SpanGroup(doc, name=self.key)