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