Initial upload: aesthetic-immich-starter
This commit is contained in:
@@ -0,0 +1,35 @@
|
||||
# Copyright 2022 The MediaPipe Authors.
|
||||
#
|
||||
# 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.
|
||||
|
||||
"""MediaPipe Tasks Text API."""
|
||||
|
||||
import mediapipe.tasks.python.text.language_detector
|
||||
import mediapipe.tasks.python.text.text_classifier
|
||||
import mediapipe.tasks.python.text.text_embedder
|
||||
|
||||
LanguageDetector = language_detector.LanguageDetector
|
||||
LanguageDetectorOptions = language_detector.LanguageDetectorOptions
|
||||
LanguageDetectorResult = language_detector.LanguageDetectorResult
|
||||
TextClassifier = text_classifier.TextClassifier
|
||||
TextClassifierOptions = text_classifier.TextClassifierOptions
|
||||
TextClassifierResult = text_classifier.TextClassifierResult
|
||||
TextEmbedder = text_embedder.TextEmbedder
|
||||
TextEmbedderOptions = text_embedder.TextEmbedderOptions
|
||||
TextEmbedderResult = text_embedder.TextEmbedderResult
|
||||
|
||||
# Remove unnecessary modules to avoid duplication in API docs.
|
||||
del mediapipe
|
||||
del language_detector
|
||||
del text_classifier
|
||||
del text_embedder
|
||||
@@ -0,0 +1,290 @@
|
||||
# Copyright 2023 The MediaPipe Authors.
|
||||
#
|
||||
# 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.
|
||||
"""MediaPipe language detector task."""
|
||||
|
||||
import ctypes
|
||||
import dataclasses
|
||||
from typing import List, Optional
|
||||
from mediapipe.tasks.python.components.containers import category as category_module
|
||||
from mediapipe.tasks.python.components.containers import classification_result as classification_result_module
|
||||
from mediapipe.tasks.python.components.processors import classifier_options as classifier_options_module
|
||||
from mediapipe.tasks.python.components.processors import classifier_options_c as classifier_options_c_module
|
||||
from mediapipe.tasks.python.core import base_options as base_options_module
|
||||
from mediapipe.tasks.python.core import base_options_c as base_options_c_module
|
||||
from mediapipe.tasks.python.core import mediapipe_c_bindings as mediapipe_c_bindings_module
|
||||
from mediapipe.tasks.python.core import mediapipe_c_utils
|
||||
from mediapipe.tasks.python.core import serial_dispatcher
|
||||
|
||||
_BaseOptions = base_options_module.BaseOptions
|
||||
Category = category_module.Category
|
||||
Classifications = classification_result_module.Classifications
|
||||
ClassifierOptions = classifier_options_module.ClassifierOptions
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class LanguageDetectorResult:
|
||||
|
||||
@dataclasses.dataclass
|
||||
class Detection:
|
||||
"""A language code and its probability."""
|
||||
|
||||
# An i18n language / locale code, e.g. "en" for English, "uz" for Uzbek,
|
||||
# "ja"-Latn for Japanese (romaji).
|
||||
language_code: str
|
||||
probability: float
|
||||
|
||||
detections: List[Detection]
|
||||
|
||||
|
||||
class LanguageDetectorPredictionC(ctypes.Structure):
|
||||
"""A language code and its probability."""
|
||||
|
||||
_fields_ = [
|
||||
("language_code", ctypes.c_char_p),
|
||||
("probability", ctypes.c_float),
|
||||
]
|
||||
|
||||
|
||||
class LanguageDetectorResultC(ctypes.Structure):
|
||||
"""Language detector output."""
|
||||
|
||||
_fields_ = [
|
||||
("predictions", ctypes.POINTER(LanguageDetectorPredictionC)),
|
||||
("predictions_count", ctypes.c_uint32),
|
||||
]
|
||||
|
||||
|
||||
def _convert_to_python_language_detector_result(
|
||||
c_result: LanguageDetectorResultC,
|
||||
) -> LanguageDetectorResult:
|
||||
"""Converts a C LanguageDetectorResult to a Python LanguageDetectorResult."""
|
||||
py_result = LanguageDetectorResult(detections=[])
|
||||
for i in range(c_result.predictions_count):
|
||||
c_prediction = c_result.predictions[i]
|
||||
py_prediction = LanguageDetectorResult.Detection(
|
||||
language_code=c_prediction.language_code.decode("utf-8"),
|
||||
probability=c_prediction.probability,
|
||||
)
|
||||
py_result.detections.append(py_prediction)
|
||||
return py_result
|
||||
|
||||
|
||||
class LanguageDetectorOptionsC(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("base_options", base_options_c_module.BaseOptionsC),
|
||||
("classifier_options", classifier_options_c_module.ClassifierOptionsC),
|
||||
]
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class LanguageDetectorOptions:
|
||||
"""Options for the language detector task.
|
||||
|
||||
Attributes:
|
||||
base_options: Base options for the language detector task.
|
||||
display_names_locale: The locale to use for display names specified through
|
||||
the TFLite Model Metadata.
|
||||
max_results: The maximum number of top-scored classification results to
|
||||
return.
|
||||
score_threshold: Overrides the ones provided in the model metadata. Results
|
||||
below this value are rejected.
|
||||
category_allowlist: Allowlist of category names. If non-empty,
|
||||
classification results whose category name is not in this set will be
|
||||
filtered out. Duplicate or unknown category names are ignored. Mutually
|
||||
exclusive with `category_denylist`.
|
||||
category_denylist: Denylist of category names. If non-empty, classification
|
||||
results whose category name is in this set will be filtered out. Duplicate
|
||||
or unknown category names are ignored. Mutually exclusive with
|
||||
`category_allowlist`.
|
||||
"""
|
||||
base_options: _BaseOptions
|
||||
display_names_locale: Optional[str] = None
|
||||
max_results: Optional[int] = None
|
||||
score_threshold: Optional[float] = None
|
||||
category_allowlist: Optional[List[str]] = None
|
||||
category_denylist: Optional[List[str]] = None
|
||||
|
||||
def to_ctypes(self) -> LanguageDetectorOptionsC:
|
||||
"""Generates a ctypes LanguageDetectorOptionsC."""
|
||||
base_options_c = self.base_options.to_ctypes()
|
||||
classifier_options_c = (
|
||||
classifier_options_c_module.convert_to_classifier_options_c(
|
||||
classifier_options_module.ClassifierOptions(
|
||||
display_names_locale=self.display_names_locale,
|
||||
max_results=self.max_results,
|
||||
score_threshold=self.score_threshold,
|
||||
category_allowlist=self.category_allowlist,
|
||||
category_denylist=self.category_denylist,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
c_options = LanguageDetectorOptionsC()
|
||||
c_options.base_options = base_options_c
|
||||
c_options.classifier_options = classifier_options_c
|
||||
return c_options
|
||||
|
||||
|
||||
_CTYPES_SIGNATURES = (
|
||||
mediapipe_c_utils.CStatusFunction(
|
||||
"MpLanguageDetectorCreate",
|
||||
(
|
||||
ctypes.POINTER(LanguageDetectorOptionsC),
|
||||
ctypes.POINTER(ctypes.c_void_p),
|
||||
),
|
||||
),
|
||||
mediapipe_c_utils.CStatusFunction(
|
||||
"MpLanguageDetectorDetect",
|
||||
(
|
||||
ctypes.c_void_p,
|
||||
ctypes.c_char_p,
|
||||
ctypes.POINTER(LanguageDetectorResultC),
|
||||
),
|
||||
),
|
||||
mediapipe_c_utils.CStatusFunction(
|
||||
"MpLanguageDetectorClose",
|
||||
(
|
||||
ctypes.c_void_p,
|
||||
),
|
||||
),
|
||||
mediapipe_c_utils.CFunction(
|
||||
"MpLanguageDetectorCloseResult",
|
||||
[ctypes.POINTER(LanguageDetectorResultC)],
|
||||
None,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class LanguageDetector:
|
||||
"""Class that predicts the language of an input text.
|
||||
|
||||
This API expects a TFLite model with TFLite Model Metadata that contains the
|
||||
mandatory (described below) input tensors, output tensor, and the language
|
||||
codes in an AssociatedFile.
|
||||
|
||||
Input tensors:
|
||||
(kTfLiteString)
|
||||
- 1 input tensor that is scalar or has shape [1] containing the input
|
||||
string.
|
||||
Output tensor:
|
||||
(kTfLiteFloat32)
|
||||
- 1 output tensor of shape`[1 x N]` where `N` is the number of languages.
|
||||
"""
|
||||
_lib: serial_dispatcher.SerialDispatcher
|
||||
_handle: ctypes.c_void_p
|
||||
|
||||
def __init__(
|
||||
self, lib: serial_dispatcher.SerialDispatcher, handle: ctypes.c_void_p
|
||||
):
|
||||
self._lib = lib
|
||||
self._detector_handle = handle
|
||||
|
||||
@classmethod
|
||||
def create_from_model_path(cls, model_path: str) -> "LanguageDetector":
|
||||
"""Creates an `LanguageDetector` object from a TensorFlow Lite model and the default `LanguageDetectorOptions`.
|
||||
|
||||
Args:
|
||||
model_path: Path to the model.
|
||||
|
||||
Returns:
|
||||
`LanguageDetector` object that's created from the model file and the
|
||||
default `LanguageDetectorOptions`.
|
||||
|
||||
Raises:
|
||||
ValueError: If failed to create `LanguageDetector` object from the
|
||||
provided
|
||||
file such as invalid file path.
|
||||
RuntimeError: If other types of error occurred.
|
||||
"""
|
||||
return cls.create_from_options(
|
||||
LanguageDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=model_path)
|
||||
)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def create_from_options(
|
||||
cls, options: LanguageDetectorOptions
|
||||
) -> "LanguageDetector":
|
||||
"""Creates the `LanguageDetector` object from language detector options.
|
||||
|
||||
Args:
|
||||
options: Options for the language detector task.
|
||||
|
||||
Returns:
|
||||
`LanguageDetector` object that's created from `options`.
|
||||
|
||||
Raises:
|
||||
ValueError: If failed to create `LanguageDetector` object from
|
||||
`LanguageDetectorOptions` such as missing the model.
|
||||
RuntimeError: If other types of error occurred.
|
||||
"""
|
||||
lib = mediapipe_c_bindings_module.load_shared_library(_CTYPES_SIGNATURES)
|
||||
|
||||
ctypes_options = options.to_ctypes()
|
||||
|
||||
detector_handle = ctypes.c_void_p()
|
||||
lib.MpLanguageDetectorCreate(
|
||||
ctypes.byref(ctypes_options),
|
||||
ctypes.byref(detector_handle),
|
||||
)
|
||||
return LanguageDetector(lib=lib, handle=detector_handle)
|
||||
|
||||
def detect(self, text: str) -> LanguageDetectorResult:
|
||||
"""Predicts the language of the input `text`.
|
||||
|
||||
Args:
|
||||
text: The input text.
|
||||
|
||||
Returns:
|
||||
A `LanguageDetectorResult` object that contains a list of languages and
|
||||
scores.
|
||||
|
||||
Raises:
|
||||
ValueError: If any of the input arguments is invalid.
|
||||
RuntimeError: If language detection failed to run.
|
||||
"""
|
||||
ctypes_result = LanguageDetectorResultC()
|
||||
|
||||
self._lib.MpLanguageDetectorDetect(
|
||||
self._detector_handle,
|
||||
text.encode("utf-8"),
|
||||
ctypes.byref(ctypes_result),
|
||||
)
|
||||
python_result = _convert_to_python_language_detector_result(ctypes_result)
|
||||
self._lib.MpLanguageDetectorCloseResult(ctypes.byref(ctypes_result))
|
||||
return python_result
|
||||
|
||||
def close(self):
|
||||
"""Shuts down the MediaPipe task instance."""
|
||||
if not self._detector_handle:
|
||||
return
|
||||
self._lib.MpLanguageDetectorClose(self._detector_handle)
|
||||
self._detector_handle = None
|
||||
self._lib.close()
|
||||
|
||||
def __enter__(self):
|
||||
"""Returns `self` upon entering the runtime context."""
|
||||
return self
|
||||
|
||||
def __exit__(self, unused_exc_type, unused_exc_value, unused_traceback):
|
||||
"""Shuts down the MediaPipe task instance on exit of the context manager.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the MediaPipe LanguageDetector task failed to close.
|
||||
"""
|
||||
self.close()
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
@@ -0,0 +1,264 @@
|
||||
# Copyright 2022 The MediaPipe Authors.
|
||||
#
|
||||
# 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.
|
||||
"""MediaPipe text classifier task."""
|
||||
|
||||
import ctypes
|
||||
import dataclasses
|
||||
from typing import List, Optional
|
||||
from mediapipe.tasks.python.components.containers import category as category_module
|
||||
from mediapipe.tasks.python.components.containers import classification_result as classification_result_module
|
||||
from mediapipe.tasks.python.components.containers import classification_result_c as classification_result_c_module
|
||||
from mediapipe.tasks.python.components.processors import classifier_options as classifier_options_module
|
||||
from mediapipe.tasks.python.components.processors import classifier_options_c as classifier_options_c_module
|
||||
from mediapipe.tasks.python.core import base_options as base_options_module
|
||||
from mediapipe.tasks.python.core import base_options_c as base_options_c_module
|
||||
from mediapipe.tasks.python.core import mediapipe_c_bindings
|
||||
from mediapipe.tasks.python.core import mediapipe_c_utils
|
||||
from mediapipe.tasks.python.core import serial_dispatcher
|
||||
|
||||
_BaseOptions = base_options_module.BaseOptions
|
||||
Category = category_module.Category
|
||||
Classifications = classification_result_module.Classifications
|
||||
TextClassifierResult = classification_result_module.ClassificationResult
|
||||
ClassifierOptions = classifier_options_module.ClassifierOptions
|
||||
|
||||
|
||||
class TextClassifierOptionsC(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("base_options", base_options_c_module.BaseOptionsC),
|
||||
("classifier_options", classifier_options_c_module.ClassifierOptionsC),
|
||||
]
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class TextClassifierOptions:
|
||||
"""Options for the text classifier task.
|
||||
|
||||
Attributes:
|
||||
base_options: Base options for the text classifier task.
|
||||
display_names_locale: The locale to use for display names specified through
|
||||
the TFLite Model Metadata.
|
||||
max_results: The maximum number of top-scored classification results to
|
||||
return.
|
||||
score_threshold: Overrides the ones provided in the model metadata. Results
|
||||
below this value are rejected.
|
||||
category_allowlist: Allowlist of category names. If non-empty,
|
||||
classification results whose category name is not in this set will be
|
||||
filtered out. Duplicate or unknown category names are ignored. Mutually
|
||||
exclusive with `category_denylist`.
|
||||
category_denylist: Denylist of category names. If non-empty, classification
|
||||
results whose category name is in this set will be filtered out. Duplicate
|
||||
or unknown category names are ignored. Mutually exclusive with
|
||||
`category_allowlist`.
|
||||
"""
|
||||
base_options: _BaseOptions
|
||||
display_names_locale: Optional[str] = None
|
||||
max_results: Optional[int] = None
|
||||
score_threshold: Optional[float] = None
|
||||
category_allowlist: Optional[List[str]] = None
|
||||
category_denylist: Optional[List[str]] = None
|
||||
|
||||
def to_ctypes(self) -> TextClassifierOptionsC:
|
||||
"""Generates a ctypes TextClassifierOptionsC."""
|
||||
base_options_c = self.base_options.to_ctypes()
|
||||
classifier_options_c = (
|
||||
classifier_options_c_module.convert_to_classifier_options_c(
|
||||
classifier_options_module.ClassifierOptions(
|
||||
display_names_locale=self.display_names_locale,
|
||||
max_results=self.max_results,
|
||||
score_threshold=self.score_threshold,
|
||||
category_allowlist=self.category_allowlist,
|
||||
category_denylist=self.category_denylist,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
c_options = TextClassifierOptionsC()
|
||||
c_options.base_options = base_options_c
|
||||
c_options.classifier_options = classifier_options_c
|
||||
return c_options
|
||||
|
||||
|
||||
_CTYPES_SIGNATURES = (
|
||||
mediapipe_c_utils.CStatusFunction(
|
||||
"MpTextClassifierCreate",
|
||||
[
|
||||
ctypes.POINTER(TextClassifierOptionsC),
|
||||
ctypes.POINTER(ctypes.c_void_p),
|
||||
],
|
||||
),
|
||||
mediapipe_c_utils.CStatusFunction(
|
||||
"MpTextClassifierClassify",
|
||||
[
|
||||
ctypes.c_void_p,
|
||||
ctypes.c_char_p,
|
||||
ctypes.POINTER(
|
||||
classification_result_c_module.ClassificationResultC
|
||||
),
|
||||
],
|
||||
),
|
||||
mediapipe_c_utils.CStatusFunction(
|
||||
"MpTextClassifierClose",
|
||||
[
|
||||
ctypes.c_void_p,
|
||||
],
|
||||
),
|
||||
mediapipe_c_utils.CFunction(
|
||||
"MpTextClassifierCloseResult",
|
||||
[ctypes.POINTER(classification_result_c_module.ClassificationResultC)],
|
||||
None,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TextClassifier:
|
||||
"""Class that performs classification on text.
|
||||
|
||||
This API expects a TFLite model with (optional) TFLite Model Metadata that
|
||||
contains the mandatory (described below) input tensors, output tensor,
|
||||
and the optional (but recommended) category labels as AssociatedFiles with
|
||||
type
|
||||
TENSOR_AXIS_LABELS per output classification tensor. Metadata is required for
|
||||
models with int32 input tensors because it contains the input process unit
|
||||
for the model's Tokenizer. No metadata is required for models with string
|
||||
input tensors.
|
||||
|
||||
Input tensors:
|
||||
(kTfLiteInt32)
|
||||
- 3 input tensors of size `[batch_size x bert_max_seq_len]` representing
|
||||
the input ids, segment ids, and mask ids
|
||||
- or 1 input tensor of size `[batch_size x max_seq_len]` representing the
|
||||
input ids
|
||||
or (kTfLiteString)
|
||||
- 1 input tensor that is shapeless or has shape [1] containing the input
|
||||
string
|
||||
At least one output tensor with:
|
||||
(kTfLiteFloat32/kBool)
|
||||
- `[1 x N]` array with `N` represents the number of categories.
|
||||
- optional (but recommended) category labels as AssociatedFiles with type
|
||||
TENSOR_AXIS_LABELS, containing one label per line. The first such
|
||||
AssociatedFile (if any) is used to fill the `category_name` field of the
|
||||
results. The `display_name` field is filled from the AssociatedFile (if
|
||||
any) whose locale matches the `display_names_locale` field of the
|
||||
`TextClassifierOptions` used at creation time ("en" by default, i.e.
|
||||
English). If none of these are available, only the `index` field of the
|
||||
results will be filled.
|
||||
"""
|
||||
_lib: serial_dispatcher.SerialDispatcher
|
||||
_handle: ctypes.c_void_p
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
lib: serial_dispatcher.SerialDispatcher,
|
||||
handle: ctypes.c_void_p,
|
||||
):
|
||||
self._lib = lib
|
||||
self._classifier_handle = handle
|
||||
|
||||
@classmethod
|
||||
def create_from_model_path(cls, model_path: str) -> "TextClassifier":
|
||||
"""Creates an `TextClassifier` object from a TensorFlow Lite model and the default `TextClassifierOptions`.
|
||||
|
||||
Args:
|
||||
model_path: Path to the model.
|
||||
|
||||
Returns:
|
||||
`TextClassifier` object that's created from the model file and the
|
||||
default `TextClassifierOptions`.
|
||||
|
||||
Raises:
|
||||
ValueError: If failed to create `TextClassifier` object from the provided
|
||||
file such as invalid file path.
|
||||
RuntimeError: If other types of error occurred.
|
||||
"""
|
||||
return cls.create_from_options(
|
||||
TextClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=model_path)
|
||||
)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def create_from_options(
|
||||
cls, options: TextClassifierOptions
|
||||
) -> "TextClassifier":
|
||||
"""Creates the `TextClassifier` object from text classifier options.
|
||||
|
||||
Args:
|
||||
options: Options for the text classifier task.
|
||||
|
||||
Returns:
|
||||
`TextClassifier` object that's created from `options`.
|
||||
|
||||
Raises:
|
||||
ValueError: If failed to create `TextClassifier` object from
|
||||
`TextClassifierOptions` such as missing the model.
|
||||
RuntimeError: If other types of error occurred.
|
||||
"""
|
||||
lib = mediapipe_c_bindings.load_shared_library(_CTYPES_SIGNATURES)
|
||||
|
||||
ctypes_options = options.to_ctypes()
|
||||
|
||||
classifier_handle = ctypes.c_void_p()
|
||||
lib.MpTextClassifierCreate(
|
||||
ctypes.byref(ctypes_options),
|
||||
ctypes.byref(classifier_handle),
|
||||
)
|
||||
return TextClassifier(lib=lib, handle=classifier_handle)
|
||||
|
||||
def classify(self, text: str) -> TextClassifierResult:
|
||||
"""Performs classification on the input `text`.
|
||||
|
||||
Args:
|
||||
text: The input text.
|
||||
|
||||
Returns:
|
||||
A `TextClassifierResult` object that contains a list of text
|
||||
classifications.
|
||||
|
||||
Raises:
|
||||
ValueError: If any of the input arguments is invalid.
|
||||
RuntimeError: If text classification failed to run.
|
||||
"""
|
||||
ctypes_result = classification_result_c_module.ClassificationResultC()
|
||||
|
||||
self._lib.MpTextClassifierClassify(
|
||||
self._classifier_handle,
|
||||
text.encode("utf-8"),
|
||||
ctypes.byref(ctypes_result),
|
||||
)
|
||||
python_result = TextClassifierResult.from_ctypes(ctypes_result)
|
||||
self._lib.MpTextClassifierCloseResult(ctypes.byref(ctypes_result))
|
||||
return python_result
|
||||
|
||||
def close(self):
|
||||
"""Shuts down the MediaPipe task instance."""
|
||||
if self._classifier_handle:
|
||||
self._lib.MpTextClassifierClose(self._classifier_handle)
|
||||
self._classifier_handle = None
|
||||
self._lib.close()
|
||||
|
||||
def __enter__(self):
|
||||
"""Returns `self` upon entering the runtime context."""
|
||||
return self
|
||||
|
||||
def __exit__(self, unused_exc_type, unused_exc_value, unused_traceback):
|
||||
"""Shuts down the MediaPipe task instance on exit of the context manager.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the MediaPipe TextClassifier task failed to close.
|
||||
"""
|
||||
self.close()
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
@@ -0,0 +1,260 @@
|
||||
# Copyright 2022 The MediaPipe Authors.
|
||||
#
|
||||
# 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.
|
||||
"""MediaPipe text embedder task."""
|
||||
import ctypes
|
||||
import dataclasses
|
||||
from typing import Optional
|
||||
|
||||
from mediapipe.tasks.python.components.containers import embedding_result as embedding_result_module
|
||||
from mediapipe.tasks.python.components.containers import embedding_result_c as embedding_result_c_module
|
||||
from mediapipe.tasks.python.components.utils import cosine_similarity
|
||||
from mediapipe.tasks.python.core import base_options as base_options_module
|
||||
from mediapipe.tasks.python.core import base_options_c as base_options_c_module
|
||||
from mediapipe.tasks.python.core import mediapipe_c_bindings
|
||||
from mediapipe.tasks.python.core import mediapipe_c_utils
|
||||
from mediapipe.tasks.python.core import serial_dispatcher
|
||||
|
||||
TextEmbedderResult = embedding_result_module.EmbeddingResult
|
||||
|
||||
|
||||
class _EmbedderOptionsC(ctypes.Structure):
|
||||
"""C struct for embedder options."""
|
||||
|
||||
_fields_ = [
|
||||
('l2_normalize', ctypes.c_bool),
|
||||
('quantize', ctypes.c_bool),
|
||||
]
|
||||
|
||||
|
||||
class _TextEmbedderOptionsC(ctypes.Structure):
|
||||
"""C struct for text embedder options."""
|
||||
|
||||
_fields_ = [
|
||||
('base_options', base_options_c_module.BaseOptionsC),
|
||||
('embedder_options', _EmbedderOptionsC),
|
||||
]
|
||||
|
||||
|
||||
_CTYPES_SIGNATURES = (
|
||||
mediapipe_c_utils.CStatusFunction(
|
||||
'MpTextEmbedderCreate',
|
||||
[
|
||||
ctypes.POINTER(_TextEmbedderOptionsC),
|
||||
ctypes.POINTER(ctypes.c_void_p),
|
||||
],
|
||||
),
|
||||
mediapipe_c_utils.CStatusFunction(
|
||||
'MpTextEmbedderEmbed',
|
||||
[
|
||||
ctypes.c_void_p,
|
||||
ctypes.c_char_p,
|
||||
ctypes.POINTER(embedding_result_c_module.EmbeddingResultC),
|
||||
],
|
||||
),
|
||||
mediapipe_c_utils.CStatusFunction(
|
||||
'MpTextEmbedderClose',
|
||||
[
|
||||
ctypes.c_void_p,
|
||||
],
|
||||
),
|
||||
mediapipe_c_utils.CFunction(
|
||||
'MpTextEmbedderCloseResult',
|
||||
[ctypes.POINTER(embedding_result_c_module.EmbeddingResultC)],
|
||||
None,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class TextEmbedderOptions:
|
||||
"""Options for the text embedder task.
|
||||
|
||||
Attributes:
|
||||
base_options: Base options for the text embedder task.
|
||||
l2_normalize: Whether to normalize the returned feature vector with L2 norm.
|
||||
Use this option only if the model does not already contain a native
|
||||
L2_NORMALIZATION TF Lite Op. In most cases, this is already the case and
|
||||
L2 norm is thus achieved through TF Lite inference.
|
||||
quantize: Whether the returned embedding should be quantized to bytes via
|
||||
scalar quantization. Embeddings are implicitly assumed to be unit-norm and
|
||||
therefore any dimension is guaranteed to have a value in [-1.0, 1.0]. Use
|
||||
the l2_normalize option if this is not the case.
|
||||
"""
|
||||
base_options: base_options_module.BaseOptions
|
||||
l2_normalize: Optional[bool] = None
|
||||
quantize: Optional[bool] = None
|
||||
|
||||
def to_ctypes(self) -> _TextEmbedderOptionsC:
|
||||
"""Generates a ctypes TextEmbedderOptionsC object."""
|
||||
base_options_c = self.base_options.to_ctypes()
|
||||
embedder_options_c = _EmbedderOptionsC(
|
||||
l2_normalize=self.l2_normalize
|
||||
if self.l2_normalize is not None
|
||||
else False,
|
||||
quantize=self.quantize if self.quantize is not None else False,
|
||||
)
|
||||
return _TextEmbedderOptionsC(
|
||||
base_options=base_options_c, embedder_options=embedder_options_c
|
||||
)
|
||||
|
||||
|
||||
class TextEmbedder:
|
||||
"""Class that performs embedding extraction on text.
|
||||
|
||||
This API expects a TFLite model with TFLite Model Metadata that contains the
|
||||
mandatory (described below) input tensors and output tensors. Metadata should
|
||||
contain the input process unit for the model's Tokenizer as well as input /
|
||||
output tensor metadata.
|
||||
|
||||
Input tensors:
|
||||
(kTfLiteInt32)
|
||||
- 3 input tensors of size `[batch_size x bert_max_seq_len]` with names
|
||||
"ids", "mask", and "segment_ids" representing the input ids, mask ids, and
|
||||
segment ids respectively.
|
||||
- or 1 input tensor of size `[batch_size x max_seq_len]` representing the
|
||||
input ids.
|
||||
|
||||
At least one output tensor with:
|
||||
(kTfLiteFloat32)
|
||||
- `N` components corresponding to the `N` dimensions of the returned
|
||||
feature vector for this output layer.
|
||||
- Either 2 or 4 dimensions, i.e. `[1 x N]` or `[1 x 1 x 1 x N]`.
|
||||
"""
|
||||
_lib: serial_dispatcher.SerialDispatcher
|
||||
_handle: ctypes.c_void_p
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
lib: serial_dispatcher.SerialDispatcher,
|
||||
handle: ctypes.c_void_p,
|
||||
):
|
||||
self._lib = lib
|
||||
self._embedder_handle = handle
|
||||
|
||||
@classmethod
|
||||
def create_from_model_path(cls, model_path: str) -> 'TextEmbedder':
|
||||
"""Creates an `TextEmbedder` object from a TensorFlow Lite model and the default `TextEmbedderOptions`.
|
||||
|
||||
Args:
|
||||
model_path: Path to the model.
|
||||
|
||||
Returns:
|
||||
`TextEmbedder` object that's created from the model file and the default
|
||||
`TextEmbedderOptions`.
|
||||
|
||||
Raises:
|
||||
ValueError: If failed to create `TextEmbedder` object from the provided
|
||||
file such as invalid file path.
|
||||
RuntimeError: If other types of error occurred.
|
||||
"""
|
||||
base_options = base_options_module.BaseOptions(model_asset_path=model_path)
|
||||
options = TextEmbedderOptions(base_options=base_options)
|
||||
return cls.create_from_options(options)
|
||||
|
||||
@classmethod
|
||||
def create_from_options(cls, options: TextEmbedderOptions) -> 'TextEmbedder':
|
||||
"""Creates the `TextEmbedder` object from text embedder options.
|
||||
|
||||
Args:
|
||||
options: Options for the text embedder task.
|
||||
|
||||
Returns:
|
||||
`TextEmbedder` object that's created from `options`.
|
||||
|
||||
Raises:
|
||||
ValueError: If failed to create `TextEmbedder` object from
|
||||
`TextEmbedderOptions` such as missing the model.
|
||||
RuntimeError: If other types of error occurred.
|
||||
"""
|
||||
lib = mediapipe_c_bindings.load_shared_library(_CTYPES_SIGNATURES)
|
||||
ctypes_options = options.to_ctypes()
|
||||
embedder_handle = ctypes.c_void_p()
|
||||
lib.MpTextEmbedderCreate(
|
||||
ctypes.byref(ctypes_options), ctypes.byref(embedder_handle)
|
||||
)
|
||||
return TextEmbedder(lib=lib, handle=embedder_handle)
|
||||
|
||||
def embed(
|
||||
self,
|
||||
text: str,
|
||||
) -> embedding_result_module.EmbeddingResult:
|
||||
"""Performs text embedding extraction on the provided text.
|
||||
|
||||
Args:
|
||||
text: The input text.
|
||||
|
||||
Returns:
|
||||
An embedding result object that contains a list of embeddings.
|
||||
|
||||
Raises:
|
||||
ValueError: If any of the input arguments is invalid.
|
||||
RuntimeError: If text embedder failed to run.
|
||||
"""
|
||||
ctypes_result = embedding_result_c_module.EmbeddingResultC()
|
||||
self._lib.MpTextEmbedderEmbed(
|
||||
self._embedder_handle,
|
||||
text.encode('utf-8'),
|
||||
ctypes.byref(ctypes_result),
|
||||
)
|
||||
python_result = (
|
||||
embedding_result_module.EmbeddingResult.from_ctypes(
|
||||
ctypes_result
|
||||
)
|
||||
)
|
||||
self._lib.MpTextEmbedderCloseResult(ctypes.byref(ctypes_result))
|
||||
return python_result
|
||||
|
||||
@classmethod
|
||||
def cosine_similarity(
|
||||
cls,
|
||||
u: embedding_result_module.Embedding,
|
||||
v: embedding_result_module.Embedding,
|
||||
) -> float:
|
||||
"""Utility function to compute cosine similarity between two embedding entries.
|
||||
|
||||
May return an InvalidArgumentError if e.g. the feature vectors are
|
||||
of different types (quantized vs. float), have different sizes, or have a
|
||||
an L2-norm of 0.
|
||||
|
||||
Args:
|
||||
u: An embedding entry.
|
||||
v: An embedding entry.
|
||||
|
||||
Returns:
|
||||
The cosine similarity for the two embeddings.
|
||||
|
||||
Raises:
|
||||
ValueError: May return an error if e.g. the feature vectors are of
|
||||
different types (quantized vs. float), have different sizes, or have
|
||||
an L2-norm of 0.
|
||||
"""
|
||||
return cosine_similarity.cosine_similarity(u, v)
|
||||
|
||||
def close(self):
|
||||
"""Shuts down the MediaPipe task instance."""
|
||||
if self._embedder_handle:
|
||||
self._lib.MpTextEmbedderClose(self._embedder_handle)
|
||||
self._embedder_handle = None
|
||||
self._lib.close()
|
||||
|
||||
def __enter__(self):
|
||||
"""Returns `self` upon entering the runtime context."""
|
||||
return self
|
||||
|
||||
def __exit__(self, unused_exc_type, unused_exc_value, unused_traceback):
|
||||
"""Shuts down the MediaPipe task instance on exit of the context manager."""
|
||||
self.close()
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
Reference in New Issue
Block a user