hmz007
36ed224bac
Signed-off-by: hmz007 <hmz007@gmail.com> Change-Id: Ib87fdbe7a0be7dc961af3500fe9ea4589a127f9f
91 lines
3.0 KiB
Java
91 lines
3.0 KiB
Java
/*
|
|
* Copyright (C) 2018 The Android Open Source Project
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
package com.android.textclassifier;
|
|
|
|
import androidx.annotation.FloatRange;
|
|
import androidx.collection.ArrayMap;
|
|
import com.google.common.base.Preconditions;
|
|
import java.util.ArrayList;
|
|
import java.util.Collections;
|
|
import java.util.List;
|
|
import java.util.Map;
|
|
|
|
/** Helper object for setting and getting entity scores for classified text. */
|
|
final class EntityConfidence {
|
|
|
|
static final EntityConfidence EMPTY = new EntityConfidence(Collections.emptyMap());
|
|
|
|
private final ArrayMap<String, Float> entityConfidence = new ArrayMap<>();
|
|
private final ArrayList<String> sortedEntities = new ArrayList<>();
|
|
|
|
/**
|
|
* Constructs an EntityConfidence from a map of entity to confidence.
|
|
*
|
|
* <p>Map entries that have 0 confidence are removed, and values greater than 1 are clamped to 1.
|
|
*
|
|
* @param source a map from entity to a confidence value in the range 0 (low confidence) to 1
|
|
* (high confidence).
|
|
*/
|
|
EntityConfidence(Map<String, Float> source) {
|
|
Preconditions.checkNotNull(source);
|
|
|
|
// Prune non-existent entities and clamp to 1.
|
|
entityConfidence.ensureCapacity(source.size());
|
|
for (Map.Entry<String, Float> it : source.entrySet()) {
|
|
if (it.getValue() <= 0) {
|
|
continue;
|
|
}
|
|
entityConfidence.put(it.getKey(), Math.min(1, it.getValue()));
|
|
}
|
|
resetSortedEntitiesFromMap();
|
|
}
|
|
|
|
/**
|
|
* Returns an immutable list of entities found in the classified text ordered from high confidence
|
|
* to low confidence.
|
|
*/
|
|
public List<String> getEntities() {
|
|
return Collections.unmodifiableList(sortedEntities);
|
|
}
|
|
|
|
/**
|
|
* Returns the confidence score for the specified entity. The value ranges from 0 (low confidence)
|
|
* to 1 (high confidence). 0 indicates that the entity was not found for the classified text.
|
|
*/
|
|
@FloatRange(from = 0.0, to = 1.0)
|
|
public float getConfidenceScore(String entity) {
|
|
return entityConfidence.getOrDefault(entity, 0f);
|
|
}
|
|
|
|
@Override
|
|
public String toString() {
|
|
return entityConfidence.toString();
|
|
}
|
|
|
|
private void resetSortedEntitiesFromMap() {
|
|
sortedEntities.clear();
|
|
sortedEntities.ensureCapacity(entityConfidence.size());
|
|
sortedEntities.addAll(entityConfidence.keySet());
|
|
sortedEntities.sort(
|
|
(e1, e2) -> {
|
|
float score1 = entityConfidence.get(e1);
|
|
float score2 = entityConfidence.get(e2);
|
|
return Float.compare(score2, score1);
|
|
});
|
|
}
|
|
}
|