EntityConfidence.java revision bbe43dfd97c01364e46df452be4c99536d64e4fb
1/* 2 * Copyright (C) 2017 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17package android.view.textclassifier; 18 19import android.annotation.FloatRange; 20import android.annotation.NonNull; 21import android.util.ArrayMap; 22 23import com.android.internal.util.Preconditions; 24 25import java.util.ArrayList; 26import java.util.Collections; 27import java.util.List; 28import java.util.Map; 29 30/** 31 * Helper object for setting and getting entity scores for classified text. 32 * 33 * @param <T> the entity type. 34 * @hide 35 */ 36final class EntityConfidence<T> { 37 38 private final ArrayMap<T, Float> mEntityConfidence = new ArrayMap<>(); 39 private final ArrayList<T> mSortedEntities = new ArrayList<>(); 40 41 EntityConfidence() {} 42 43 EntityConfidence(@NonNull EntityConfidence<T> source) { 44 Preconditions.checkNotNull(source); 45 mEntityConfidence.putAll(source.mEntityConfidence); 46 mSortedEntities.addAll(source.mSortedEntities); 47 } 48 49 /** 50 * Constructs an EntityConfidence from a map of entity to confidence. 51 * 52 * Map entries that have 0 confidence are removed, and values greater than 1 are clamped to 1. 53 * 54 * @param source a map from entity to a confidence value in the range 0 (low confidence) to 55 * 1 (high confidence). 56 */ 57 EntityConfidence(@NonNull Map<T, Float> source) { 58 Preconditions.checkNotNull(source); 59 60 // Prune non-existent entities and clamp to 1. 61 mEntityConfidence.ensureCapacity(source.size()); 62 for (Map.Entry<T, Float> it : source.entrySet()) { 63 if (it.getValue() <= 0) continue; 64 mEntityConfidence.put(it.getKey(), Math.min(1, it.getValue())); 65 } 66 67 // Create a list of entities sorted by decreasing confidence for getEntities(). 68 mSortedEntities.ensureCapacity(mEntityConfidence.size()); 69 mSortedEntities.addAll(mEntityConfidence.keySet()); 70 mSortedEntities.sort((e1, e2) -> { 71 float score1 = mEntityConfidence.get(e1); 72 float score2 = mEntityConfidence.get(e2); 73 return Float.compare(score2, score1); 74 }); 75 } 76 77 /** 78 * Returns an immutable list of entities found in the classified text ordered from 79 * high confidence to low confidence. 80 */ 81 @NonNull 82 public List<T> getEntities() { 83 return Collections.unmodifiableList(mSortedEntities); 84 } 85 86 /** 87 * Returns the confidence score for the specified entity. The value ranges from 88 * 0 (low confidence) to 1 (high confidence). 0 indicates that the entity was not found for the 89 * classified text. 90 */ 91 @FloatRange(from = 0.0, to = 1.0) 92 public float getConfidenceScore(T entity) { 93 if (mEntityConfidence.containsKey(entity)) { 94 return mEntityConfidence.get(entity); 95 } 96 return 0; 97 } 98 99 @Override 100 public String toString() { 101 return mEntityConfidence.toString(); 102 } 103} 104