Limit the number of entity spans detected by NER, fix `get_token_spans()`

app.py

@@ -80,6 +80,9 @@ def get_token_spans(tokenizer, text: str) -> list[tuple[int, int]]:
     end = 0
     for token in tokenizer.tokenize(text):
         token = token.removeprefix("##")
+        if token not in text:  # e.g., token == "[UNK]"
+            continue
+
         start = text.index(token, end)
         end = start + len(token)
         token_spans.append((start, end))
@@ -164,6 +167,7 @@ def get_topk_entities_from_texts(
         model_outputs = model(**tokenized_examples)
         token_spans = get_token_spans(tokenizer, text)
         entity_spans = get_predicted_entity_spans(model_outputs.ner_logits[0], token_spans, entity_span_sensitivity)
+        entity_spans = entity_spans[:tokenizer.max_entity_length]
         batch_entity_spans.append(entity_spans)
 
         tokenized_examples = tokenizer(text, entity_spans=entity_spans or None, truncation=True, return_tensors="pt")