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