EntityConfidence.java revision c08637c95be5e31268baf9b434a2d710423181f5
1c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus/* 2c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * Copyright (C) 2018 The Android Open Source Project 3c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * 4c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * Licensed under the Apache License, Version 2.0 (the "License"); 5c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * you may not use this file except in compliance with the License. 6c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * You may obtain a copy of the License at 7c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * 8c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * http://www.apache.org/licenses/LICENSE-2.0 9c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * 10c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * Unless required by applicable law or agreed to in writing, software 11c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * distributed under the License is distributed on an "AS IS" BASIS, 12c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * See the License for the specific language governing permissions and 14c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * limitations under the License. 15c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus */ 16c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 17c08637c95be5e31268baf9b434a2d710423181f5Jan Althauspackage androidx.textclassifier; 18c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 19c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport android.os.Parcel; 20c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport android.os.Parcelable; 21c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport android.support.annotation.FloatRange; 22c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport android.support.annotation.NonNull; 23c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport android.support.annotation.RestrictTo; 24c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport android.support.v4.util.ArrayMap; 25c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport android.support.v4.util.Preconditions; 26c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport android.support.v4.util.SimpleArrayMap; 27c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 28c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport java.util.ArrayList; 29c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport java.util.Collections; 30c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport java.util.Comparator; 31c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport java.util.List; 32c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport java.util.Map; 33c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 34c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus/** 35c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * Helper object for setting and getting entity scores for classified text. 36c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * 37c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * @hide 38c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus */ 39c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus@RestrictTo(RestrictTo.Scope.LIBRARY) 40c08637c95be5e31268baf9b434a2d710423181f5Jan Althausfinal class EntityConfidence implements Parcelable { 41c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 42c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus private final ArrayMap<String, Float> mEntityConfidence = new ArrayMap<>(); 43c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus private final ArrayList<String> mSortedEntities = new ArrayList<>(); 44c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 45c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus EntityConfidence() {} 46c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 47c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus EntityConfidence(@NonNull EntityConfidence source) { 48c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus Preconditions.checkNotNull(source); 49c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mEntityConfidence.putAll((SimpleArrayMap<String, Float>) source.mEntityConfidence); 50c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mSortedEntities.addAll(source.mSortedEntities); 51c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 52c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 53c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus /** 54c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * Constructs an EntityConfidence from a map of entity to confidence. 55c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * 56c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * Map entries that have 0 confidence are removed, and values greater than 1 are clamped to 1. 57c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * 58c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * @param source a map from entity to a confidence value in the range 0 (low confidence) to 59c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * 1 (high confidence). 60c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus */ 61c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus EntityConfidence(@NonNull Map<String, Float> source) { 62c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus Preconditions.checkNotNull(source); 63c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 64c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus // Prune non-existent entities and clamp to 1. 65c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mEntityConfidence.ensureCapacity(source.size()); 66c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus for (Map.Entry<String, Float> it : source.entrySet()) { 67c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus if (it.getValue() <= 0) continue; 68c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mEntityConfidence.put(it.getKey(), Math.min(1, it.getValue())); 69c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 70c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus resetSortedEntitiesFromMap(); 71c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 72c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 73c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus /** 74c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * Returns an immutable list of entities found in the classified text ordered from 75c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * high confidence to low confidence. 76c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus */ 77c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @NonNull 78c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public List<String> getEntities() { 79c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return Collections.unmodifiableList(mSortedEntities); 80c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 81c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 82c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus /** 83c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * Returns the confidence score for the specified entity. The value ranges from 84c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * 0 (low confidence) to 1 (high confidence). 0 indicates that the entity was not found for the 85c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * classified text. 86c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus */ 87c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @FloatRange(from = 0.0, to = 1.0) 88c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public float getConfidenceScore(String entity) { 89c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus if (mEntityConfidence.containsKey(entity)) { 90c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return mEntityConfidence.get(entity); 91c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 92c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return 0; 93c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 94c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 95c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @Override 96c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public String toString() { 97c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return mEntityConfidence.toString(); 98c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 99c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 100c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @Override 101c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public int describeContents() { 102c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return 0; 103c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 104c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 105c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @Override 106c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public void writeToParcel(Parcel dest, int flags) { 107c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus dest.writeInt(mEntityConfidence.size()); 108c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus for (Map.Entry<String, Float> entry : mEntityConfidence.entrySet()) { 109c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus dest.writeString(entry.getKey()); 110c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus dest.writeFloat(entry.getValue()); 111c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 112c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 113c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 114c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public static final Parcelable.Creator<EntityConfidence> CREATOR = 115c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus new Parcelable.Creator<EntityConfidence>() { 116c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @Override 117c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public EntityConfidence createFromParcel(Parcel in) { 118c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return new EntityConfidence(in); 119c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 120c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 121c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @Override 122c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public EntityConfidence[] newArray(int size) { 123c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return new EntityConfidence[size]; 124c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 125c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus }; 126c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 127c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus private EntityConfidence(Parcel in) { 128c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus final int numEntities = in.readInt(); 129c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mEntityConfidence.ensureCapacity(numEntities); 130c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus for (int i = 0; i < numEntities; ++i) { 131c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mEntityConfidence.put(in.readString(), in.readFloat()); 132c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 133c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus resetSortedEntitiesFromMap(); 134c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 135c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 136c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus private void resetSortedEntitiesFromMap() { 137c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mSortedEntities.clear(); 138c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mSortedEntities.ensureCapacity(mEntityConfidence.size()); 139c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mSortedEntities.addAll(mEntityConfidence.keySet()); 140c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus Collections.sort(mSortedEntities, new EntityConfidenceComparator()); 141c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 142c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 143c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus /** Helper to sort entities according to their confidence. */ 144c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus private class EntityConfidenceComparator implements Comparator<String> { 145c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @Override 146c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public int compare(String e1, String e2) { 147c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus float score1 = mEntityConfidence.get(e1); 148c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus float score2 = mEntityConfidence.get(e2); 149c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return Float.compare(score2, score1); 150c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 151c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 152c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus} 153