blob: bf74f8fcd29cf2af8ab82daa50edbf3645907241 [file] [log] [blame]
Honglin Yuf33dce32019-12-05 15:10:39 +11001// Copyright 2020 The Chromium OS 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
5#ifndef ML_TEXT_CLASSIFIER_IMPL_H_
6#define ML_TEXT_CLASSIFIER_IMPL_H_
7
8#include <list>
9#include <map>
10#include <memory>
11#include <string>
12
13#include <annotator/annotator.h>
14#include <base/macros.h>
Andrew Moylanb481af72020-07-09 15:22:00 +100015#include <mojo/public/cpp/bindings/pending_receiver.h>
16#include <mojo/public/cpp/bindings/receiver.h>
Honglin Yuf33dce32019-12-05 15:10:39 +110017#include <tensorflow/lite/model.h>
18#include <utils/memory/mmap.h>
19
20#include "ml/graph_executor_impl.h"
21#include "ml/mojom/model.mojom.h"
22#include "ml/mojom/text_classifier.mojom.h"
23
24namespace ml {
25
26class TextClassifierImpl
27 : public chromeos::machine_learning::mojom::TextClassifier {
28 public:
Andrew Moylan79b34a42020-07-08 11:13:11 +100029 // Interface to create new `TextClassifierImpl` object. This function will
Andrew Moylanb481af72020-07-09 15:22:00 +100030 // automatically achieve strong binding.The model object will be deleted when
31 // the corresponding mojo connection is closed.
Honglin Yuf33dce32019-12-05 15:10:39 +110032 // Will return false if it fails to create the annotator object, otherwise
33 // return true.
34 static bool Create(
Andrew Moylanb481af72020-07-09 15:22:00 +100035 mojo::PendingReceiver<chromeos::machine_learning::mojom::TextClassifier>
36 receiver);
Honglin Yuf33dce32019-12-05 15:10:39 +110037
38 private:
Andrew Moylan79b34a42020-07-08 11:13:11 +100039 // A private constructor, call `TextClassifierImpl::Create` to create new
Honglin Yuf33dce32019-12-05 15:10:39 +110040 // objects.
41 explicit TextClassifierImpl(
Andrew Moylanb481af72020-07-09 15:22:00 +100042 std::unique_ptr<libtextclassifier3::ScopedMmap>* mmap,
Honglin Yuc5100022020-07-09 11:54:27 +100043 const std::string& langid_model_path,
Andrew Moylanb481af72020-07-09 15:22:00 +100044 mojo::PendingReceiver<chromeos::machine_learning::mojom::TextClassifier>
45 receiver);
Honglin Yuf33dce32019-12-05 15:10:39 +110046
Andrew Moylanb481af72020-07-09 15:22:00 +100047 void SetDisconnectionHandler(base::Closure disconnect_handler);
Honglin Yuf33dce32019-12-05 15:10:39 +110048
49 // chromeos::machine_learning::mojom::TextClassifier:
50 void Annotate(
51 chromeos::machine_learning::mojom::TextAnnotationRequestPtr request,
52 AnnotateCallback callback) override;
53
54 // chromeos::machine_learning::mojom::TextClassifier:
55 void SuggestSelection(
56 chromeos::machine_learning::mojom::TextSuggestSelectionRequestPtr request,
57 SuggestSelectionCallback callback) override;
58
Honglin Yuc5100022020-07-09 11:54:27 +100059 // chromeos::machine_learning::mojom::TextClassifier:
60 void FindLanguages(const std::string& text,
61 FindLanguagesCallback callback) override;
62
Honglin Yuf33dce32019-12-05 15:10:39 +110063 std::unique_ptr<libtextclassifier3::Annotator> annotator_;
64
Honglin Yuc5100022020-07-09 11:54:27 +100065 std::unique_ptr<libtextclassifier3::mobile::lang_id::LangId>
66 language_identifier_;
67
Andrew Moylanb481af72020-07-09 15:22:00 +100068 mojo::Receiver<chromeos::machine_learning::mojom::TextClassifier> receiver_;
Honglin Yuf33dce32019-12-05 15:10:39 +110069
70 DISALLOW_COPY_AND_ASSIGN(TextClassifierImpl);
71};
72
73} // namespace ml
74
75#endif // ML_TEXT_CLASSIFIER_IMPL_H_