small changes to entity mention ner confidence setup

stanfordnlp · Sep 20, 2018 · 1e934ee · 1e934ee
1 parent 7212edc
commit 1e934ee
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 24 deletions.
diff --git a/src/edu/stanford/nlp/ie/NERClassifierCombiner.java b/src/edu/stanford/nlp/ie/NERClassifierCombiner.java
@@ -290,10 +290,10 @@ public boolean usesSUTime() {
   private static <INN extends CoreMap> void copyAnswerFieldsToNERField(List<INN> l) {
     for (INN m: l) {
       m.set(CoreAnnotations.NamedEntityTagAnnotation.class, m.get(CoreAnnotations.AnswerAnnotation.class));
-      HashMap<String,Double> labelToProb = new HashMap<String,Double>();
-      labelToProb.put(
-          m.get(CoreAnnotations.NamedEntityTagAnnotation.class), m.get(CoreAnnotations.AnswerProbAnnotation.class));
-      m.set(CoreAnnotations.NamedEntityTagProbAnnotation.class, labelToProb);
+      Map<String,Double> labelToProb =
+          Collections.singletonMap(m.get(CoreAnnotations.NamedEntityTagAnnotation.class),
+              m.get(CoreAnnotations.AnswerProbAnnotation.class));
+      m.set(CoreAnnotations.NamedEntityTagProbsAnnotation.class, labelToProb);
     }
   }
 

diff --git a/src/edu/stanford/nlp/ling/CoreAnnotations.java b/src/edu/stanford/nlp/ling/CoreAnnotations.java
@@ -92,10 +92,10 @@ public Class<String> getType() {
   /**
    * Label and probability pair representing the coarse grained label and probability
    */
-  public static class NamedEntityTagProbAnnotation implements CoreAnnotation<HashMap<String,Double>> {
+  public static class NamedEntityTagProbsAnnotation implements CoreAnnotation<Map<String,Double>> {
     @Override
-    public Class<HashMap<String,Double>> getType() {
-      return ErasureUtils.uncheckedCast(HashMap.class);
+    public Class<Map<String,Double>> getType() {
+      return ErasureUtils.uncheckedCast(Map.class);
     }
   }
 

diff --git a/src/edu/stanford/nlp/pipeline/EntityMentionsAnnotator.java b/src/edu/stanford/nlp/pipeline/EntityMentionsAnnotator.java
@@ -221,7 +221,7 @@ public static HashMap<String,Double> determineEntityMentionConfidences(CoreMap e
     // get a list of labels that have probability values from the first token
     Set<String> labelsWithProbs =
         entityMention.get(CoreAnnotations.TokensAnnotation.class).get(0).get(
-            CoreAnnotations.NamedEntityTagProbAnnotation.class).keySet();
+            CoreAnnotations.NamedEntityTagProbsAnnotation.class).keySet();
     // build the label values hash map for the entity mention
     HashMap<String,Double> entityLabelProbVals = new HashMap<>();
     // initialize to 1.1
@@ -230,8 +230,8 @@ public static HashMap<String,Double> determineEntityMentionConfidences(CoreMap e
     }
     // go through each token, see if you can find a smaller prob value for that label
     for (CoreLabel token : entityMention.get(CoreAnnotations.TokensAnnotation.class)) {
-      HashMap<String,Double> labelProbsForToken =
-          token.get(CoreAnnotations.NamedEntityTagProbAnnotation.class);
+      Map<String,Double> labelProbsForToken =
+          token.get(CoreAnnotations.NamedEntityTagProbsAnnotation.class);
       for (String label : labelProbsForToken.keySet()) {
         if (labelProbsForToken.get(label) < entityLabelProbVals.get(label))
           entityLabelProbVals.put(label, labelProbsForToken.get(label));
@@ -331,7 +331,7 @@ public void annotate(Annotation annotation) {
     for (CoreMap entityMention : allEntityMentions) {
       HashMap<String,Double> entityMentionLabelProbVals =
           determineEntityMentionConfidences(entityMention);
-      entityMention.set(CoreAnnotations.NamedEntityTagProbAnnotation.class, entityMentionLabelProbVals);
+      entityMention.set(CoreAnnotations.NamedEntityTagProbsAnnotation.class, entityMentionLabelProbVals);
     }
 
     annotation.set(mentionsCoreAnnotationClass, allEntityMentions);

diff --git a/src/edu/stanford/nlp/pipeline/NERCombinerAnnotator.java b/src/edu/stanford/nlp/pipeline/NERCombinerAnnotator.java
@@ -371,10 +371,9 @@ public void annotate(Annotation annotation) {
 
     // set confidence for anything not already set to n.e. tag, -1.0
     for (CoreLabel token : annotation.get(CoreAnnotations.TokensAnnotation.class)) {
-      if (token.get(CoreAnnotations.NamedEntityTagProbAnnotation.class) == null) {
-        HashMap<String,Double> labelToProb = new HashMap<>();
-        labelToProb.put(token.ner(), -1.0);
-        token.set(CoreAnnotations.NamedEntityTagProbAnnotation.class, labelToProb);
+      if (token.get(CoreAnnotations.NamedEntityTagProbsAnnotation.class) == null) {
+        Map<String,Double> labelToProb = Collections.singletonMap(token.ner(), -1.0);
+        token.set(CoreAnnotations.NamedEntityTagProbsAnnotation.class, labelToProb);
       }
     }
 
@@ -411,13 +410,13 @@ public void doOneSentence(Annotation annotation, CoreMap sentence) {
         // add the named entity tag to each token
         String neTag = output.get(i).get(CoreAnnotations.NamedEntityTagAnnotation.class);
         String normNeTag = output.get(i).get(CoreAnnotations.NormalizedNamedEntityTagAnnotation.class);
-        HashMap<String,Double> neTagProbMap = output.get(i).get(CoreAnnotations.NamedEntityTagProbAnnotation.class);
+        Map<String,Double> neTagProbMap = output.get(i).get(CoreAnnotations.NamedEntityTagProbsAnnotation.class);
         if (language.equals(LanguageInfo.HumanLanguage.SPANISH)) {
           neTag = spanishToEnglishTag(neTag);
           normNeTag = spanishToEnglishTag(normNeTag);
         }
         tokens.get(i).setNER(neTag);
-        tokens.get(i).set(CoreAnnotations.NamedEntityTagProbAnnotation.class, neTagProbMap);
+        tokens.get(i).set(CoreAnnotations.NamedEntityTagProbsAnnotation.class, neTagProbMap);
         tokens.get(i).set(CoreAnnotations.CoarseNamedEntityTagAnnotation.class, neTag);
         if (normNeTag != null) tokens.get(i).set(CoreAnnotations.NormalizedNamedEntityTagAnnotation.class, normNeTag);
         NumberSequenceClassifier.transferAnnotations(output.get(i), tokens.get(i));

diff --git a/src/edu/stanford/nlp/pipeline/ProtobufAnnotationSerializer.java b/src/edu/stanford/nlp/pipeline/ProtobufAnnotationSerializer.java
@@ -297,18 +297,18 @@ protected CoreNLPProtos.Token.Builder toProtoBuilder(CoreLabel coreLabel, Set<Cl
       builder.setFineGrainedNER(coreLabel.get(FineGrainedNamedEntityTagAnnotation.class));
       keysToSerialize.remove(FineGrainedNamedEntityTagAnnotation.class);
     }
-    if (keySet.contains(NamedEntityTagProbAnnotation.class)) {
+    if (keySet.contains(NamedEntityTagProbsAnnotation.class)) {
       // in case of empty label prob list, add string "empty"
       // this is to differentiate between null and an empty hash map
-      if (coreLabel.get(NamedEntityTagProbAnnotation.class).keySet().size() == 0) {
+      if (coreLabel.get(NamedEntityTagProbsAnnotation.class).keySet().size() == 0) {
         builder.addNerLabelProbs("empty");
       } else {
-        for (String labelWithProb : coreLabel.get(NamedEntityTagProbAnnotation.class).keySet()) {
-          Double labelProb = coreLabel.get(NamedEntityTagProbAnnotation.class).get(labelWithProb);
+        for (String labelWithProb : coreLabel.get(NamedEntityTagProbsAnnotation.class).keySet()) {
+          Double labelProb = coreLabel.get(NamedEntityTagProbsAnnotation.class).get(labelWithProb);
           builder.addNerLabelProbs(labelWithProb+"="+labelProb);
         }
       }
-      keysToSerialize.remove(NamedEntityTagProbAnnotation.class);
+      keysToSerialize.remove(NamedEntityTagProbsAnnotation.class);
     }
     if (keySet.contains(CharacterOffsetBeginAnnotation.class)) { builder.setBeginChar(coreLabel.beginPosition()); keysToSerialize.remove(CharacterOffsetBeginAnnotation.class); }
     if (keySet.contains(CharacterOffsetEndAnnotation.class)) { builder.setEndChar(coreLabel.endPosition()); keysToSerialize.remove(CharacterOffsetEndAnnotation.class); }
@@ -1400,7 +1400,7 @@ public CoreLabel fromProto(CoreNLPProtos.Token proto) {
           nerLabelProbs.put(labelAndProb[0], labelProbDouble);
         }
       }
-      word.set(NamedEntityTagProbAnnotation.class, nerLabelProbs);
+      word.set(NamedEntityTagProbsAnnotation.class, nerLabelProbs);
     }
 
     // Return
@@ -1745,7 +1745,7 @@ public Annotation fromProto(CoreNLPProtos.Document proto) {
           // set entity mention label prob vals
           HashMap<String,Double> nerLabelConfidences =
               EntityMentionsAnnotator.determineEntityMentionConfidences(entityMention);
-          entityMention.set(CoreAnnotations.NamedEntityTagProbAnnotation.class, nerLabelConfidences);
+          entityMention.set(CoreAnnotations.NamedEntityTagProbsAnnotation.class, nerLabelConfidences);
         }
         if (sentence.getHasEntityMentionsAnnotation()) {
           map.set(CoreAnnotations.MentionsAnnotation.class, mentions);