bert_model_handler.h 1.3 KB

12345678910111213141516171819202122232425262728293031323334
  1. // Copyright 2021 The Chromium Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license that can be
  3. // found in the LICENSE file.
  4. #ifndef COMPONENTS_OPTIMIZATION_GUIDE_CORE_BERT_MODEL_HANDLER_H_
  5. #define COMPONENTS_OPTIMIZATION_GUIDE_CORE_BERT_MODEL_HANDLER_H_
  6. #include "components/optimization_guide/core/model_handler.h"
  7. #include "third_party/tflite_support/src/tensorflow_lite_support/cc/task/core/category.h"
  8. namespace optimization_guide {
  9. // An implementation of a ModelHandler that executes BERT models.
  10. //
  11. // Note that sentencepiece tokenizers are not supported by Chromium's copy of
  12. // the TFLite Support library.
  13. class BertModelHandler
  14. : public ModelHandler<std::vector<tflite::task::core::Category>,
  15. const std::string&> {
  16. public:
  17. BertModelHandler(
  18. OptimizationGuideModelProvider* model_provider,
  19. scoped_refptr<base::SequencedTaskRunner> background_task_runner,
  20. proto::OptimizationTarget optimization_target,
  21. const absl::optional<proto::Any>& model_metadata);
  22. ~BertModelHandler() override;
  23. BertModelHandler(const BertModelHandler&) = delete;
  24. BertModelHandler& operator=(const BertModelHandler&) = delete;
  25. };
  26. } // namespace optimization_guide
  27. #endif // COMPONENTS_OPTIMIZATION_GUIDE_CORE_BERT_MODEL_HANDLER_H_