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