/*
|
* 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.
|
*/
|
|
// Inference code for the text classification model.
|
|
#ifndef LIBTEXTCLASSIFIER_ANNOTATOR_ANNOTATOR_H_
|
#define LIBTEXTCLASSIFIER_ANNOTATOR_ANNOTATOR_H_
|
|
#include <memory>
|
#include <set>
|
#include <string>
|
#include <unordered_set>
|
#include <vector>
|
|
#include "annotator/contact/contact-engine.h"
|
#include "annotator/datetime/parser.h"
|
#include "annotator/duration/duration.h"
|
#include "annotator/feature-processor.h"
|
#include "annotator/installed_app/installed-app-engine.h"
|
#include "annotator/knowledge/knowledge-engine.h"
|
#include "annotator/model-executor.h"
|
#include "annotator/model_generated.h"
|
#include "annotator/number/number.h"
|
#include "annotator/strip-unpaired-brackets.h"
|
#include "annotator/types.h"
|
#include "annotator/zlib-utils.h"
|
#include "utils/flatbuffers.h"
|
#include "utils/i18n/locale.h"
|
#include "utils/memory/mmap.h"
|
#include "utils/utf8/unilib.h"
|
#include "utils/zlib/zlib.h"
|
|
namespace libtextclassifier3 {
|
|
// Aliases for long enum values.
|
const AnnotationUsecase ANNOTATION_USECASE_SMART =
|
AnnotationUsecase_ANNOTATION_USECASE_SMART;
|
const AnnotationUsecase ANNOTATION_USECASE_RAW =
|
AnnotationUsecase_ANNOTATION_USECASE_RAW;
|
|
struct SelectionOptions {
|
// Comma-separated list of locale specification for the input text (BCP 47
|
// tags).
|
std::string locales;
|
|
// Comma-separated list of BCP 47 language tags.
|
std::string detected_text_language_tags;
|
|
// Tailors the output annotations according to the specified use-case.
|
AnnotationUsecase annotation_usecase = ANNOTATION_USECASE_SMART;
|
|
bool operator==(const SelectionOptions& other) const {
|
return this->locales == other.locales &&
|
this->annotation_usecase == other.annotation_usecase &&
|
this->detected_text_language_tags ==
|
other.detected_text_language_tags;
|
}
|
};
|
|
struct ClassificationOptions {
|
// For parsing relative datetimes, the reference now time against which the
|
// relative datetimes get resolved.
|
// UTC milliseconds since epoch.
|
int64 reference_time_ms_utc = 0;
|
|
// Timezone in which the input text was written (format as accepted by ICU).
|
std::string reference_timezone;
|
|
// Comma-separated list of locale specification for the input text (BCP 47
|
// tags).
|
std::string locales;
|
|
// Comma-separated list of language tags.
|
std::string detected_text_language_tags;
|
|
// Tailors the output annotations according to the specified use-case.
|
AnnotationUsecase annotation_usecase = ANNOTATION_USECASE_SMART;
|
|
bool operator==(const ClassificationOptions& other) const {
|
return this->reference_time_ms_utc == other.reference_time_ms_utc &&
|
this->reference_timezone == other.reference_timezone &&
|
this->locales == other.locales &&
|
this->detected_text_language_tags ==
|
other.detected_text_language_tags &&
|
this->annotation_usecase == other.annotation_usecase;
|
}
|
};
|
|
struct AnnotationOptions {
|
// For parsing relative datetimes, the reference now time against which the
|
// relative datetimes get resolved.
|
// UTC milliseconds since epoch.
|
int64 reference_time_ms_utc = 0;
|
|
// Timezone in which the input text was written (format as accepted by ICU).
|
std::string reference_timezone;
|
|
// Comma-separated list of locale specification for the input text (BCP 47
|
// tags).
|
std::string locales;
|
|
// Comma-separated list of language tags.
|
std::string detected_text_language_tags;
|
|
// List of entity types that should be used for annotation.
|
std::unordered_set<std::string> entity_types;
|
|
// If true, serialized_entity_data in the results is populated."
|
bool is_serialized_entity_data_enabled = false;
|
|
// Tailors the output annotations according to the specified use-case.
|
AnnotationUsecase annotation_usecase = ANNOTATION_USECASE_SMART;
|
|
bool operator==(const AnnotationOptions& other) const {
|
return this->reference_time_ms_utc == other.reference_time_ms_utc &&
|
this->reference_timezone == other.reference_timezone &&
|
this->locales == other.locales &&
|
this->detected_text_language_tags ==
|
other.detected_text_language_tags &&
|
this->annotation_usecase == other.annotation_usecase &&
|
this->is_serialized_entity_data_enabled ==
|
other.is_serialized_entity_data_enabled;
|
}
|
};
|
|
// Holds TFLite interpreters for selection and classification models.
|
// NOTE: his class is not thread-safe, thus should NOT be re-used across
|
// threads.
|
class InterpreterManager {
|
public:
|
// The constructor can be called with nullptr for any of the executors, and is
|
// a defined behavior, as long as the corresponding *Interpreter() method is
|
// not called when the executor is null.
|
InterpreterManager(const ModelExecutor* selection_executor,
|
const ModelExecutor* classification_executor)
|
: selection_executor_(selection_executor),
|
classification_executor_(classification_executor) {}
|
|
// Gets or creates and caches an interpreter for the selection model.
|
tflite::Interpreter* SelectionInterpreter();
|
|
// Gets or creates and caches an interpreter for the classification model.
|
tflite::Interpreter* ClassificationInterpreter();
|
|
private:
|
const ModelExecutor* selection_executor_;
|
const ModelExecutor* classification_executor_;
|
|
std::unique_ptr<tflite::Interpreter> selection_interpreter_;
|
std::unique_ptr<tflite::Interpreter> classification_interpreter_;
|
};
|
|
// Stores entity types enabled for annotation, and provides operator() for
|
// checking whether a given entity type is enabled.
|
class EnabledEntityTypes {
|
public:
|
explicit EnabledEntityTypes(
|
const std::unordered_set<std::string>& entity_types)
|
: entity_types_(entity_types) {}
|
|
bool operator()(const std::string& entity_type) const {
|
return entity_types_.empty() ||
|
entity_types_.find(entity_type) != entity_types_.cend();
|
}
|
|
private:
|
const std::unordered_set<std::string>& entity_types_;
|
};
|
|
// A text processing model that provides text classification, annotation,
|
// selection suggestion for various types.
|
// NOTE: This class is not thread-safe.
|
class Annotator {
|
public:
|
static std::unique_ptr<Annotator> FromUnownedBuffer(
|
const char* buffer, int size, const UniLib* unilib = nullptr,
|
const CalendarLib* calendarlib = nullptr);
|
// Takes ownership of the mmap.
|
static std::unique_ptr<Annotator> FromScopedMmap(
|
std::unique_ptr<ScopedMmap>* mmap, const UniLib* unilib = nullptr,
|
const CalendarLib* calendarlib = nullptr);
|
static std::unique_ptr<Annotator> FromScopedMmap(
|
std::unique_ptr<ScopedMmap>* mmap, std::unique_ptr<UniLib> unilib,
|
std::unique_ptr<CalendarLib> calendarlib);
|
static std::unique_ptr<Annotator> FromFileDescriptor(
|
int fd, int offset, int size, const UniLib* unilib = nullptr,
|
const CalendarLib* calendarlib = nullptr);
|
static std::unique_ptr<Annotator> FromFileDescriptor(
|
int fd, int offset, int size, std::unique_ptr<UniLib> unilib,
|
std::unique_ptr<CalendarLib> calendarlib);
|
static std::unique_ptr<Annotator> FromFileDescriptor(
|
int fd, const UniLib* unilib = nullptr,
|
const CalendarLib* calendarlib = nullptr);
|
static std::unique_ptr<Annotator> FromFileDescriptor(
|
int fd, std::unique_ptr<UniLib> unilib,
|
std::unique_ptr<CalendarLib> calendarlib);
|
static std::unique_ptr<Annotator> FromPath(
|
const std::string& path, const UniLib* unilib = nullptr,
|
const CalendarLib* calendarlib = nullptr);
|
static std::unique_ptr<Annotator> FromPath(
|
const std::string& path, std::unique_ptr<UniLib> unilib,
|
std::unique_ptr<CalendarLib> calendarlib);
|
|
// Returns true if the model is ready for use.
|
bool IsInitialized() { return initialized_; }
|
|
// Initializes the knowledge engine with the given config.
|
bool InitializeKnowledgeEngine(const std::string& serialized_config);
|
|
// Initializes the contact engine with the given config.
|
bool InitializeContactEngine(const std::string& serialized_config);
|
|
// Initializes the installed app engine with the given config.
|
bool InitializeInstalledAppEngine(const std::string& serialized_config);
|
|
// Runs inference for given a context and current selection (i.e. index
|
// of the first and one past last selected characters (utf8 codepoint
|
// offsets)). Returns the indices (utf8 codepoint offsets) of the selection
|
// beginning character and one past selection end character.
|
// Returns the original click_indices if an error occurs.
|
// NOTE: The selection indices are passed in and returned in terms of
|
// UTF8 codepoints (not bytes).
|
// Requires that the model is a smart selection model.
|
CodepointSpan SuggestSelection(
|
const std::string& context, CodepointSpan click_indices,
|
const SelectionOptions& options = SelectionOptions()) const;
|
|
// Classifies the selected text given the context string.
|
// Returns an empty result if an error occurs.
|
std::vector<ClassificationResult> ClassifyText(
|
const std::string& context, CodepointSpan selection_indices,
|
const ClassificationOptions& options = ClassificationOptions()) const;
|
|
// Annotates given input text. The annotations are sorted by their position
|
// in the context string and exclude spans classified as 'other'.
|
std::vector<AnnotatedSpan> Annotate(
|
const std::string& context,
|
const AnnotationOptions& options = AnnotationOptions()) const;
|
|
// Looks up a knowledge entity by its id. If successful, populates the
|
// serialized knowledge result and returns true.
|
bool LookUpKnowledgeEntity(const std::string& id,
|
std::string* serialized_knowledge_result) const;
|
|
const Model* model() const;
|
const reflection::Schema* entity_data_schema() const;
|
|
// Exposes the feature processor for tests and evaluations.
|
const FeatureProcessor* SelectionFeatureProcessorForTests() const;
|
const FeatureProcessor* ClassificationFeatureProcessorForTests() const;
|
|
// Exposes the date time parser for tests and evaluations.
|
const DatetimeParser* DatetimeParserForTests() const;
|
|
static const std::string& kPhoneCollection;
|
static const std::string& kAddressCollection;
|
static const std::string& kDateCollection;
|
static const std::string& kUrlCollection;
|
static const std::string& kEmailCollection;
|
|
protected:
|
struct ScoredChunk {
|
TokenSpan token_span;
|
float score;
|
};
|
|
// Constructs and initializes text classifier from given model.
|
// Takes ownership of 'mmap', and thus owns the buffer that backs 'model'.
|
Annotator(std::unique_ptr<ScopedMmap>* mmap, const Model* model,
|
const UniLib* unilib, const CalendarLib* calendarlib);
|
Annotator(std::unique_ptr<ScopedMmap>* mmap, const Model* model,
|
std::unique_ptr<UniLib> unilib,
|
std::unique_ptr<CalendarLib> calendarlib);
|
|
// Constructs, validates and initializes text classifier from given model.
|
// Does not own the buffer that backs 'model'.
|
Annotator(const Model* model, const UniLib* unilib,
|
const CalendarLib* calendarlib);
|
|
// Checks that model contains all required fields, and initializes internal
|
// datastructures.
|
void ValidateAndInitialize();
|
|
// Initializes regular expressions for the regex model.
|
bool InitializeRegexModel(ZlibDecompressor* decompressor);
|
|
// Resolves conflicts in the list of candidates by removing some overlapping
|
// ones. Returns indices of the surviving ones.
|
// NOTE: Assumes that the candidates are sorted according to their position in
|
// the span.
|
bool ResolveConflicts(const std::vector<AnnotatedSpan>& candidates,
|
const std::string& context,
|
const std::vector<Token>& cached_tokens,
|
const std::vector<Locale>& detected_text_language_tags,
|
AnnotationUsecase annotation_usecase,
|
InterpreterManager* interpreter_manager,
|
std::vector<int>* result) const;
|
|
// Resolves one conflict between candidates on indices 'start_index'
|
// (inclusive) and 'end_index' (exclusive). Assigns the winning candidate
|
// indices to 'chosen_indices'. Returns false if a problem arises.
|
bool ResolveConflict(const std::string& context,
|
const std::vector<Token>& cached_tokens,
|
const std::vector<AnnotatedSpan>& candidates,
|
const std::vector<Locale>& detected_text_language_tags,
|
int start_index, int end_index,
|
AnnotationUsecase annotation_usecase,
|
InterpreterManager* interpreter_manager,
|
std::vector<int>* chosen_indices) const;
|
|
// Gets selection candidates from the ML model.
|
// Provides the tokens produced during tokenization of the context string for
|
// reuse.
|
bool ModelSuggestSelection(
|
const UnicodeText& context_unicode, CodepointSpan click_indices,
|
const std::vector<Locale>& detected_text_language_tags,
|
InterpreterManager* interpreter_manager, std::vector<Token>* tokens,
|
std::vector<AnnotatedSpan>* result) const;
|
|
// Classifies the selected text given the context string with the
|
// classification model.
|
// Returns true if no error occurred.
|
bool ModelClassifyText(
|
const std::string& context, const std::vector<Token>& cached_tokens,
|
const std::vector<Locale>& locales, CodepointSpan selection_indices,
|
InterpreterManager* interpreter_manager,
|
FeatureProcessor::EmbeddingCache* embedding_cache,
|
std::vector<ClassificationResult>* classification_results,
|
std::vector<Token>* tokens) const;
|
|
// Same as above but doesn't output tokens.
|
bool ModelClassifyText(
|
const std::string& context, const std::vector<Token>& cached_tokens,
|
const std::vector<Locale>& detected_text_language_tags,
|
CodepointSpan selection_indices, InterpreterManager* interpreter_manager,
|
FeatureProcessor::EmbeddingCache* embedding_cache,
|
std::vector<ClassificationResult>* classification_results) const;
|
|
// Same as above but doesn't take cached tokens and doesn't output tokens.
|
bool ModelClassifyText(
|
const std::string& context,
|
const std::vector<Locale>& detected_text_language_tags,
|
CodepointSpan selection_indices, InterpreterManager* interpreter_manager,
|
FeatureProcessor::EmbeddingCache* embedding_cache,
|
std::vector<ClassificationResult>* classification_results) const;
|
|
// Returns a relative token span that represents how many tokens on the left
|
// from the selection and right from the selection are needed for the
|
// classifier input.
|
TokenSpan ClassifyTextUpperBoundNeededTokens() const;
|
|
// Classifies the selected text with the regular expressions models.
|
// Returns true if no error happened, false otherwise.
|
bool RegexClassifyText(
|
const std::string& context, CodepointSpan selection_indices,
|
std::vector<ClassificationResult>* classification_result) const;
|
|
// Classifies the selected text with the date time model.
|
// Returns true if no error happened, false otherwise.
|
bool DatetimeClassifyText(
|
const std::string& context, CodepointSpan selection_indices,
|
const ClassificationOptions& options,
|
std::vector<ClassificationResult>* classification_results) const;
|
|
// Chunks given input text with the selection model and classifies the spans
|
// with the classification model.
|
// The annotations are sorted by their position in the context string and
|
// exclude spans classified as 'other'.
|
// Provides the tokens produced during tokenization of the context string for
|
// reuse.
|
bool ModelAnnotate(const std::string& context,
|
const std::vector<Locale>& detected_text_language_tags,
|
InterpreterManager* interpreter_manager,
|
std::vector<Token>* tokens,
|
std::vector<AnnotatedSpan>* result) const;
|
|
// Groups the tokens into chunks. A chunk is a token span that should be the
|
// suggested selection when any of its contained tokens is clicked. The chunks
|
// are non-overlapping and are sorted by their position in the context string.
|
// "num_tokens" is the total number of tokens available (as this method does
|
// not need the actual vector of tokens).
|
// "span_of_interest" is a span of all the tokens that could be clicked.
|
// The resulting chunks all have to overlap with it and they cover this span
|
// completely. The first and last chunk might extend beyond it.
|
// The chunks vector is cleared before filling.
|
bool ModelChunk(int num_tokens, const TokenSpan& span_of_interest,
|
tflite::Interpreter* selection_interpreter,
|
const CachedFeatures& cached_features,
|
std::vector<TokenSpan>* chunks) const;
|
|
// A helper method for ModelChunk(). It generates scored chunk candidates for
|
// a click context model.
|
// NOTE: The returned chunks can (and most likely do) overlap.
|
bool ModelClickContextScoreChunks(
|
int num_tokens, const TokenSpan& span_of_interest,
|
const CachedFeatures& cached_features,
|
tflite::Interpreter* selection_interpreter,
|
std::vector<ScoredChunk>* scored_chunks) const;
|
|
// A helper method for ModelChunk(). It generates scored chunk candidates for
|
// a bounds-sensitive model.
|
// NOTE: The returned chunks can (and most likely do) overlap.
|
bool ModelBoundsSensitiveScoreChunks(
|
int num_tokens, const TokenSpan& span_of_interest,
|
const TokenSpan& inference_span, const CachedFeatures& cached_features,
|
tflite::Interpreter* selection_interpreter,
|
std::vector<ScoredChunk>* scored_chunks) const;
|
|
// Produces chunks isolated by a set of regular expressions.
|
bool RegexChunk(const UnicodeText& context_unicode,
|
const std::vector<int>& rules,
|
std::vector<AnnotatedSpan>* result,
|
bool is_serialized_entity_data_enabled) const;
|
|
// Produces chunks from the datetime parser.
|
bool DatetimeChunk(const UnicodeText& context_unicode,
|
int64 reference_time_ms_utc,
|
const std::string& reference_timezone,
|
const std::string& locales, ModeFlag mode,
|
AnnotationUsecase annotation_usecase,
|
bool is_serialized_entity_data_enabled,
|
std::vector<AnnotatedSpan>* result) const;
|
|
// Returns whether a classification should be filtered.
|
bool FilteredForAnnotation(const AnnotatedSpan& span) const;
|
bool FilteredForClassification(
|
const ClassificationResult& classification) const;
|
bool FilteredForSelection(const AnnotatedSpan& span) const;
|
|
// Computes the selection boundaries from a regular expression match.
|
CodepointSpan ComputeSelectionBoundaries(
|
const UniLib::RegexMatcher* match,
|
const RegexModel_::Pattern* config) const;
|
|
// Returns whether a regex pattern provides entity data from a match.
|
bool HasEntityData(const RegexModel_::Pattern* pattern) const;
|
|
// Constructs and serializes entity data from regex matches.
|
bool SerializedEntityDataFromRegexMatch(
|
const RegexModel_::Pattern* pattern, UniLib::RegexMatcher* matcher,
|
std::string* serialized_entity_data) const;
|
|
// Verifies a regex match and returns true if verification was successful.
|
bool VerifyRegexMatchCandidate(
|
const std::string& context,
|
const VerificationOptions* verification_options, const std::string& match,
|
const UniLib::RegexMatcher* matcher) const;
|
|
const Model* model_;
|
|
std::unique_ptr<const ModelExecutor> selection_executor_;
|
std::unique_ptr<const ModelExecutor> classification_executor_;
|
std::unique_ptr<const EmbeddingExecutor> embedding_executor_;
|
|
std::unique_ptr<const FeatureProcessor> selection_feature_processor_;
|
std::unique_ptr<const FeatureProcessor> classification_feature_processor_;
|
|
std::unique_ptr<const DatetimeParser> datetime_parser_;
|
|
private:
|
struct CompiledRegexPattern {
|
const RegexModel_::Pattern* config;
|
std::unique_ptr<UniLib::RegexPattern> pattern;
|
};
|
|
// Removes annotations the entity type of which is not in the set of enabled
|
// entity types.
|
void RemoveNotEnabledEntityTypes(
|
const EnabledEntityTypes& is_entity_type_enabled,
|
std::vector<AnnotatedSpan>* annotated_spans) const;
|
|
std::unique_ptr<ScopedMmap> mmap_;
|
bool initialized_ = false;
|
bool enabled_for_annotation_ = false;
|
bool enabled_for_classification_ = false;
|
bool enabled_for_selection_ = false;
|
std::unordered_set<std::string> filtered_collections_annotation_;
|
std::unordered_set<std::string> filtered_collections_classification_;
|
std::unordered_set<std::string> filtered_collections_selection_;
|
|
std::vector<CompiledRegexPattern> regex_patterns_;
|
|
// Indices into regex_patterns_ for the different modes.
|
std::vector<int> annotation_regex_patterns_, classification_regex_patterns_,
|
selection_regex_patterns_;
|
|
std::unique_ptr<UniLib> owned_unilib_;
|
const UniLib* unilib_;
|
std::unique_ptr<CalendarLib> owned_calendarlib_;
|
const CalendarLib* calendarlib_;
|
|
std::unique_ptr<const KnowledgeEngine> knowledge_engine_;
|
std::unique_ptr<const ContactEngine> contact_engine_;
|
std::unique_ptr<const InstalledAppEngine> installed_app_engine_;
|
std::unique_ptr<const NumberAnnotator> number_annotator_;
|
std::unique_ptr<const DurationAnnotator> duration_annotator_;
|
|
// Builder for creating extra data.
|
const reflection::Schema* entity_data_schema_;
|
std::unique_ptr<ReflectiveFlatbufferBuilder> entity_data_builder_;
|
|
// Locales for which the entire model triggers.
|
std::vector<Locale> model_triggering_locales_;
|
|
// Locales for which the ML model triggers.
|
std::vector<Locale> ml_model_triggering_locales_;
|
|
// Locales that the dictionary classification support.
|
std::vector<Locale> dictionary_locales_;
|
};
|
|
namespace internal {
|
|
// Helper function, which if the initial 'span' contains only white-spaces,
|
// moves the selection to a single-codepoint selection on the left side
|
// of this block of white-space.
|
CodepointSpan SnapLeftIfWhitespaceSelection(CodepointSpan span,
|
const UnicodeText& context_unicode,
|
const UniLib& unilib);
|
|
// Copies tokens from 'cached_tokens' that are
|
// 'tokens_around_selection_to_copy' (on the left, and right) tokens distant
|
// from the tokens that correspond to 'selection_indices'.
|
std::vector<Token> CopyCachedTokens(const std::vector<Token>& cached_tokens,
|
CodepointSpan selection_indices,
|
TokenSpan tokens_around_selection_to_copy);
|
} // namespace internal
|
|
// Interprets the buffer as a Model flatbuffer and returns it for reading.
|
const Model* ViewModel(const void* buffer, int size);
|
|
// Opens model from given path and runs a function, passing the loaded Model
|
// flatbuffer as an argument.
|
//
|
// This is mainly useful if we don't want to pay the cost for the model
|
// initialization because we'll be only reading some flatbuffer values from the
|
// file.
|
template <typename ReturnType, typename Func>
|
ReturnType VisitAnnotatorModel(const std::string& path, Func function) {
|
ScopedMmap mmap(path);
|
if (!mmap.handle().ok()) {
|
function(/*model=*/nullptr);
|
}
|
const Model* model =
|
ViewModel(mmap.handle().start(), mmap.handle().num_bytes());
|
return function(model);
|
}
|
|
} // namespace libtextclassifier3
|
|
#endif // LIBTEXTCLASSIFIER_ANNOTATOR_ANNOTATOR_H_
|
