EntityConfidence.java revision a6096f6c4c56c4937b52efb593d4333db301d6ad
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; 21 22import com.android.internal.util.Preconditions; 23 24import java.util.ArrayList; 25import java.util.Collections; 26import java.util.Comparator; 27import java.util.HashMap; 28import java.util.List; 29import java.util.Map; 30 31/** 32 * Helper object for setting and getting entity scores for classified text. 33 * 34 * @param <T> the entity type. 35 * @hide 36 */ 37final class EntityConfidence<T> { 38 39 private final Map<T, Float> mEntityConfidence = new HashMap<>(); 40 41 private final Comparator<T> mEntityComparator = (e1, e2) -> { 42 float score1 = mEntityConfidence.get(e1); 43 float score2 = mEntityConfidence.get(e2); 44 if (score1 > score2) { 45 return -1; 46 } 47 if (score1 < score2) { 48 return 1; 49 } 50 return 0; 51 }; 52 53 EntityConfidence() {} 54 55 EntityConfidence(@NonNull EntityConfidence<T> source) { 56 Preconditions.checkNotNull(source); 57 mEntityConfidence.putAll(source.mEntityConfidence); 58 } 59 60 /** 61 * Sets an entity type for the classified text and assigns a confidence score. 62 * 63 * @param confidenceScore a value from 0 (low confidence) to 1 (high confidence). 64 * 0 implies the entity does not exist for the classified text. 65 * Values greater than 1 are clamped to 1. 66 */ 67 public void setEntityType( 68 @NonNull T type, @FloatRange(from = 0.0, to = 1.0) float confidenceScore) { 69 Preconditions.checkNotNull(type); 70 if (confidenceScore > 0) { 71 mEntityConfidence.put(type, Math.min(1, confidenceScore)); 72 } else { 73 mEntityConfidence.remove(type); 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 List<T> entities = new ArrayList<>(mEntityConfidence.size()); 84 entities.addAll(mEntityConfidence.keySet()); 85 entities.sort(mEntityComparator); 86 return Collections.unmodifiableList(entities); 87 } 88 89 /** 90 * Returns the confidence score for the specified entity. The value ranges from 91 * 0 (low confidence) to 1 (high confidence). 0 indicates that the entity was not found for the 92 * classified text. 93 */ 94 @FloatRange(from = 0.0, to = 1.0) 95 public float getConfidenceScore(T entity) { 96 if (mEntityConfidence.containsKey(entity)) { 97 return mEntityConfidence.get(entity); 98 } 99 return 0; 100 } 101 102 @Override 103 public String toString() { 104 return mEntityConfidence.toString(); 105 } 106} 107