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 17ddf3bdd8f2a3afe5d8fd041f542ad4005a2ac869Aurimas Liutikaspackage androidx.textclassifier; 18c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 19c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport android.os.Parcel; 20c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport android.os.Parcelable; 21c65e03f6e729f8c5e5a766a1076ee3531dbdc0b2Aurimas Liutikas 22ac5fe7c617c66850fff75a9fce9979c6e5674b0fAurimas Liutikasimport androidx.annotation.FloatRange; 23ac5fe7c617c66850fff75a9fce9979c6e5674b0fAurimas Liutikasimport androidx.annotation.NonNull; 24ac5fe7c617c66850fff75a9fce9979c6e5674b0fAurimas Liutikasimport androidx.annotation.RestrictTo; 25ac5fe7c617c66850fff75a9fce9979c6e5674b0fAurimas Liutikasimport androidx.collection.ArrayMap; 26ac5fe7c617c66850fff75a9fce9979c6e5674b0fAurimas Liutikasimport androidx.collection.SimpleArrayMap; 27c65e03f6e729f8c5e5a766a1076ee3531dbdc0b2Aurimas Liutikasimport androidx.core.util.Preconditions; 28c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 29c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport java.util.ArrayList; 30c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport java.util.Collections; 31c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport java.util.Comparator; 32c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport java.util.List; 33c08637c95be5e31268baf9b434a2d710423181f5Jan Althausimport java.util.Map; 34c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 35c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus/** 36c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * Helper object for setting and getting entity scores for classified text. 37c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * 38c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * @hide 39c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus */ 40c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus@RestrictTo(RestrictTo.Scope.LIBRARY) 41c08637c95be5e31268baf9b434a2d710423181f5Jan Althausfinal class EntityConfidence implements Parcelable { 42c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 43c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus private final ArrayMap<String, Float> mEntityConfidence = new ArrayMap<>(); 44c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus private final ArrayList<String> mSortedEntities = new ArrayList<>(); 45c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 46c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus EntityConfidence() {} 47c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 48c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus EntityConfidence(@NonNull EntityConfidence source) { 49c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus Preconditions.checkNotNull(source); 50c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mEntityConfidence.putAll((SimpleArrayMap<String, Float>) source.mEntityConfidence); 51c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mSortedEntities.addAll(source.mSortedEntities); 52c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 53c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 54c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus /** 55c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * Constructs an EntityConfidence from a map of entity to confidence. 56c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * 57c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * Map entries that have 0 confidence are removed, and values greater than 1 are clamped to 1. 58c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * 59c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * @param source a map from entity to a confidence value in the range 0 (low confidence) to 60c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * 1 (high confidence). 61c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus */ 62c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus EntityConfidence(@NonNull Map<String, Float> source) { 63c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus Preconditions.checkNotNull(source); 64c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 65c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus // Prune non-existent entities and clamp to 1. 66c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mEntityConfidence.ensureCapacity(source.size()); 67c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus for (Map.Entry<String, Float> it : source.entrySet()) { 68c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus if (it.getValue() <= 0) continue; 69c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mEntityConfidence.put(it.getKey(), Math.min(1, it.getValue())); 70c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 71c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus resetSortedEntitiesFromMap(); 72c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 73c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 74c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus /** 75c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * Returns an immutable list of entities found in the classified text ordered from 76c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * high confidence to low confidence. 77c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus */ 78c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @NonNull 79c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public List<String> getEntities() { 80c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return Collections.unmodifiableList(mSortedEntities); 81c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 82c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 83c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus /** 84c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * Returns the confidence score for the specified entity. The value ranges from 85c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * 0 (low confidence) to 1 (high confidence). 0 indicates that the entity was not found for the 86c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus * classified text. 87c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus */ 88c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @FloatRange(from = 0.0, to = 1.0) 89c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public float getConfidenceScore(String entity) { 90c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus if (mEntityConfidence.containsKey(entity)) { 91c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return mEntityConfidence.get(entity); 92c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 93c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return 0; 94c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 95c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 96c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @Override 97c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public String toString() { 98c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return mEntityConfidence.toString(); 99c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 100c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 101c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @Override 102c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public int describeContents() { 103c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return 0; 104c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 105c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 106c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @Override 107c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public void writeToParcel(Parcel dest, int flags) { 108c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus dest.writeInt(mEntityConfidence.size()); 109c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus for (Map.Entry<String, Float> entry : mEntityConfidence.entrySet()) { 110c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus dest.writeString(entry.getKey()); 111c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus dest.writeFloat(entry.getValue()); 112c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 113c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 114c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 115c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public static final Parcelable.Creator<EntityConfidence> CREATOR = 116c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus new Parcelable.Creator<EntityConfidence>() { 117c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @Override 118c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public EntityConfidence createFromParcel(Parcel in) { 119c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return new EntityConfidence(in); 120c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 121c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 122c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @Override 123c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public EntityConfidence[] newArray(int size) { 124c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return new EntityConfidence[size]; 125c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 126c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus }; 127c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 128c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus private EntityConfidence(Parcel in) { 129c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus final int numEntities = in.readInt(); 130c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mEntityConfidence.ensureCapacity(numEntities); 131c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus for (int i = 0; i < numEntities; ++i) { 132c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mEntityConfidence.put(in.readString(), in.readFloat()); 133c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 134c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus resetSortedEntitiesFromMap(); 135c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 136c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 137c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus private void resetSortedEntitiesFromMap() { 138c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mSortedEntities.clear(); 139c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mSortedEntities.ensureCapacity(mEntityConfidence.size()); 140c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus mSortedEntities.addAll(mEntityConfidence.keySet()); 141c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus Collections.sort(mSortedEntities, new EntityConfidenceComparator()); 142c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 143c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus 144c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus /** Helper to sort entities according to their confidence. */ 145c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus private class EntityConfidenceComparator implements Comparator<String> { 146c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus @Override 147c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus public int compare(String e1, String e2) { 148c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus float score1 = mEntityConfidence.get(e1); 149c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus float score2 = mEntityConfidence.get(e2); 150c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus return Float.compare(score2, score1); 151c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 152c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus } 153c08637c95be5e31268baf9b434a2d710423181f5Jan Althaus} 154