krickert commented on code in PR #1086:
URL: https://github.com/apache/opennlp/pull/1086#discussion_r3434471569
##########
opennlp-core/opennlp-ml/opennlp-dl/src/main/java/opennlp/dl/namefinder/NameFinderDL.java:
##########
@@ -144,246 +152,334 @@ private static InferenceOptions
validateConstructorArguments(
@Override
public Span[] find(String[] input) {
- final List<Span> spans = new LinkedList<>();
+ final List<Span> spans = new ArrayList<>();
// Join the tokens here because they will be tokenized using Wordpiece
during inference.
final String text = String.join(" ", input);
- final String[] sentences = sentenceDetector.sentDetect(text);
+ // sentPosDetect (not sentDetect) so each sentence's offset in the full
text is known.
+ final Span[] sentenceSpans = sentenceDetector.sentPosDetect(text);
+
+ for (final Span sentenceSpan : sentenceSpans) {
- for (String sentence : sentences) {
+ // Floor the character cursor at this sentence's start, then thread it
forward across the
+ // sentence's chunks so a repeated surface form is located at its next
occurrence. Flooring
+ // per sentence keeps an entity from being matched against an identical
surface form in an
+ // earlier sentence -- even one that produced no spans, which would
otherwise leave the
+ // cursor behind and mis-locate the match.
+ int searchStart = sentenceSpan.getStart();
// The WordPiece tokenized text. This changes the spacing in the text.
- final List<Tokens> wordpieceTokens = tokenize(sentence);
+ final List<Tokens> wordpieceTokens =
tokenize(sentenceSpan.getCoveredText(text).toString());
for (final Tokens tokens : wordpieceTokens) {
+ final List<Span> decoded =
+ decodeSpans(text, tokens.tokens(), infer(tokens), ids2Labels,
searchStart);
+ spans.addAll(decoded);
+ if (!decoded.isEmpty()) {
+ searchStart = decoded.get(decoded.size() - 1).getEnd();
+ }
+ }
- try {
-
- // The inputs to the ONNX model.
- final Map<String, OnnxTensor> inputs = new HashMap<>();
-
- final float[][][] v;
- try {
- inputs.put(INPUT_IDS, OnnxTensor.createTensor(env,
LongBuffer.wrap(tokens.ids()),
- new long[] {1, tokens.ids().length}));
-
- if (includeAttentionMask) {
- inputs.put(ATTENTION_MASK, OnnxTensor.createTensor(env,
- LongBuffer.wrap(tokens.mask()), new long[] {1,
tokens.mask().length}));
- }
-
- if (includeTokenTypeIds) {
- inputs.put(TOKEN_TYPE_IDS, OnnxTensor.createTensor(env,
- LongBuffer.wrap(tokens.types()), new long[] {1,
tokens.types().length}));
- }
-
- // The outputs from the model.
- try (OrtSession.Result result = session.run(inputs)) {
- // getValue() copies the tensor into Java arrays, so the result
can be closed safely.
- v = (float[][][]) result.get(0).getValue();
- }
- } finally {
- inputs.values().forEach(OnnxTensor::close);
- }
-
- // Find consecutive B-PER and I-PER labels and combine the spans
where necessary.
- // There are also B-LOC and I-LOC tags for locations that might be
useful at some point.
+ }
- // Keep track of where the last span was so when there are
multiple/duplicate
- // spans we can get the next one instead of the first one each time.
- int characterStart = 0;
+ return spans.toArray(new Span[0]);
- final String[] toks = tokens.tokens();
+ }
- // We are looping over the vector for each word,
- // finding the index of the array that has the maximum value,
- // and then finding the token classification that corresponds to
that index.
- for (int x = 0; x < v[0].length; x++) {
+ /**
+ * Runs the model on one token window and returns the per-token label score
rows. A failure
+ * executing the model (an {@link OrtException} or any runtime fault) is
surfaced as an
+ * {@link IllegalStateException} (cause preserved); an unexpected output
shape is its own loud
+ * failure. This mirrors the fail-loud contract of the sibling {@code
DocumentCategorizerDL}.
+ *
+ * @param tokens The tokens for one chunk to run inference on.
+ * @return The {@code [token][label]} score matrix for the chunk.
+ */
+ private float[][] infer(final Tokens tokens) {
- final float[] arr = v[0][x];
- final int maxIndex = maxIndex(arr);
- final String label = ids2Labels.get(maxIndex);
+ final Map<String, OnnxTensor> inputs = new HashMap<>();
+ final Object output;
+ try {
+ inputs.put(INPUT_IDS, OnnxTensor.createTensor(env,
LongBuffer.wrap(tokens.ids()),
+ new long[] {1, tokens.ids().length}));
- // TODO: Need to make sure this value is between 0 and 1?
- // Can we do thresholding without it between 0 and 1?
- final double confidence = arr[maxIndex]; // / 10;
+ if (includeAttentionMask) {
+ inputs.put(ATTENTION_MASK, OnnxTensor.createTensor(env,
+ LongBuffer.wrap(tokens.mask()), new long[] {1,
tokens.mask().length}));
+ }
- // Is this is the start of a person entity.
- if (B_PER.equals(label)) {
+ if (includeTokenTypeIds) {
+ inputs.put(TOKEN_TYPE_IDS, OnnxTensor.createTensor(env,
+ LongBuffer.wrap(tokens.types()), new long[] {1,
tokens.types().length}));
+ }
- String spanText;
+ // getValue() copies the tensor into Java arrays, so the result can be
closed safely.
+ try (OrtSession.Result result = session.run(inputs)) {
+ output = result.get(0).getValue();
+ }
+ } catch (OrtException | RuntimeException ex) {
+ throw new IllegalStateException("Unable to perform name finder
inference", ex);
+ } finally {
+ inputs.values().forEach(OnnxTensor::close);
+ }
- // Find the end index of the span in the array (where the label
is not I-PER).
- final SpanEnd spanEnd = findSpanEnd(v, x, ids2Labels, toks);
+ // The model returns one score row per token, batched:
float[batch][token][label]. Any other
+ // shape (or an empty batch) is a model-contract violation, surfaced on
its own rather than as
+ // "inference failed".
+ if (output instanceof float[][][] v && v.length > 0) {
+ return v[0];
+ }
+ throw new IllegalStateException("Unexpected model output type: "
+ + (output == null ? "null" : output.getClass().getName()));
+ }
- // If the end is -1 it means this is a single-span token.
- // If the end is != -1 it means this is a multi-span token.
- if (spanEnd.index() != -1) {
+ @Override
+ public void clearAdaptiveData() {
+ // No use in this implementation.
+ }
- final StringBuilder sb = new StringBuilder();
+ /**
+ * Decodes spans beginning the character search at the start of {@code
text}. Equivalent to
+ * {@link #decodeSpans(String, String[], float[][], Map, int)} with {@code
searchStart == 0}.
+ *
+ * @param text The original text passed to the model.
+ * @param tokens The WordPiece tokens produced for the text.
+ * @param tokenLabelScores The per-token label scores returned by the model.
+ * @param id2Labels The mapping from model output indexes to BIO labels.
+ * @return The decoded spans.
+ */
+ static List<Span> decodeSpans(String text, String[] tokens, float[][]
tokenLabelScores,
+ Map<Integer, String> id2Labels) {
+ return decodeSpans(text, tokens, tokenLabelScores, id2Labels, 0);
+ }
- // We have to concatenate the tokens.
- // Add each token in the array and separate them with a space.
- // We'll separate each with a single space because later we'll
find the original span
- // in the text and ignore spacing between individual tokens in
findByRegex().
- int end = spanEnd.index();
- for (int i = x; i <= end; i++) {
+ /**
+ * Converts model token classifications into character spans in the original
input text.
+ *
+ * <p>The ONNX model returns one score vector for each WordPiece token. This
method applies
+ * BIO decoding, reconstructs WordPiece fragments, and then resolves the
reconstructed text
+ * against the original sentence so that {@link
Span#getCoveredText(CharSequence)} works with
+ * the caller's input.</p>
+ *
+ * @param text The original text passed to the model.
+ * @param tokens The WordPiece tokens produced for the text.
+ * @param tokenLabelScores The per-token label scores returned by the model.
+ * @param id2Labels The mapping from model output indexes to BIO labels.
+ * @param searchStart The character offset in {@code text} to begin locating
spans from. Threading
+ * a monotonic cursor across the chunks and sentences of a single {@link
#find(String[])} call
+ * keeps a repeated entity surface form from being emitted twice at the
same first occurrence.
+ * @return The decoded spans.
+ */
+ static List<Span> decodeSpans(String text, String[] tokens, float[][]
tokenLabelScores,
+ Map<Integer, String> id2Labels, int
searchStart) {
- // If the next token starts with ##, combine it with this
token.
- if (toks[i + 1].startsWith(CHARS_TO_REPLACE)) {
+ if (tokens.length != tokenLabelScores.length) {
+ throw new IllegalArgumentException("The number of tokens (" +
tokens.length
+ + ") must match the number of model output rows (" +
tokenLabelScores.length + ").");
+ }
- sb.append(toks[i]).append(toks[i +
1].replace(CHARS_TO_REPLACE, ""));
+ final List<Span> spans = new ArrayList<>();
- // Append a space unless the next (next) token starts with
##.
- if (!toks[i + 2].startsWith(CHARS_TO_REPLACE)) {
- sb.append(" ");
- }
+ int characterStart = searchStart;
- // Skip the next token since we just included it in this
iteration.
- i++;
+ for (int x = 0; x < tokenLabelScores.length; x++) {
+ final LabelPrediction prediction = predictLabel(tokenLabelScores[x],
id2Labels);
+ if (!isBeginLabel(prediction.label())) {
+ continue;
+ }
- } else {
+ final String entityType =
prediction.label().substring(BEGIN_PREFIX.length());
+ final EntityPrediction entity = findEntityEnd(tokenLabelScores, x,
id2Labels,
+ entityType, prediction.probability());
+ final String spanText = buildSpanText(tokens, x, entity.endIndex());
- sb.append(toks[i].replace(CHARS_TO_REPLACE, ""));
+ if (spanText.isBlank()) {
+ x = entity.endIndex();
+ continue;
+ }
- // Append a space unless the next token is a period.
- if (!".".equals(toks[i + 1])) {
- sb.append(" ");
- }
+ final SpanMatch match = findByRegex(text, spanText, characterStart);
+ if (match.start() != -1) {
+ spans.add(new Span(match.start(), match.end(), entityType,
entity.probability()));
+ characterStart = match.end();
+ }
- }
+ x = entity.endIndex();
+ }
- }
+ return spans;
- // This is the text of the span. We use the whole original
input text and not one
- // of the splits. This gives us accurate character positions.
- spanText = findByRegex(text, sb.toString().trim()).trim();
+ }
- } else {
+ private static EntityPrediction findEntityEnd(float[][] tokenLabelScores,
int startIndex,
+ Map<Integer, String> id2Labels,
+ String entityType,
+ double startProbability) {
- // This is a single-token span so there is nothing else to do
except grab the token.
- spanText = toks[x];
+ final String insideLabel = INSIDE_PREFIX + entityType;
+ int endIndex = startIndex;
+ double probability = startProbability;
- }
+ for (int x = startIndex + 1; x < tokenLabelScores.length; x++) {
+ final LabelPrediction prediction = predictLabel(tokenLabelScores[x],
id2Labels);
+ if (!insideLabel.equals(prediction.label())) {
+ break;
+ }
+ endIndex = x;
+ probability = Math.min(probability, prediction.probability());
+ }
- if (!SEPARATOR.equals(spanText)) {
+ return new EntityPrediction(endIndex, probability);
- spanText = spanText.replace(CHARS_TO_REPLACE, "");
+ }
- // This ignores other potential matches in the same sentence
- // by only taking the first occurrence.
- characterStart = text.indexOf(spanText, characterStart);
+ private static boolean isBeginLabel(String label) {
+ return label.startsWith(BEGIN_PREFIX) && label.length() >
BEGIN_PREFIX.length();
+ }
- // TODO: This check should not be needed because the span was
found.
- // If we aren't finding it now it's because there's a
whitespace difference.
- if (characterStart != -1) {
+ private static LabelPrediction predictLabel(float[] scores, Map<Integer,
String> id2Labels) {
- final int characterEnd = characterStart + spanText.length();
+ final int labelIndex = maxIndex(scores);
+ final String label = id2Labels.get(labelIndex);
+ if (label == null) {
+ throw new IllegalArgumentException("No label is configured for model
output index "
Review Comment:
Done - it now throws `IllegalStateException`, naming the offending model
output index. We deliberately kept this fail-loud rather than degrading the
token to `O`: an unmapped argmax index means `ids2Labels` is not exhaustive
over the model's outputs, which is a misconfiguration. The constructors now
document that `ids2Labels` must be exhaustive. (This is the fail-loud +
documented-contract option from @rzo1's note, rather than silent degradation.)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
