google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Utilities for handling model inputs."""
17

18
from __future__ import absolute_import
19
from __future__ import division
20
from __future__ import print_function
21

22
import numpy as np
23
import tensorflow.compat.v1 as tf
24

25

26
def get_shared_feature_specs(config):
27
  """Non-task-specific model inputs."""
28
  return [
29
      FeatureSpec("input_ids", [config.max_seq_length]),
30
      FeatureSpec("input_mask", [config.max_seq_length]),
31
      FeatureSpec("segment_ids", [config.max_seq_length]),
32
      FeatureSpec("task_id", []),
33
  ]
34

35

36
class FeatureSpec(object):
37
  """Defines a feature passed as input to the model."""
38

39
  def __init__(self, name, shape, default_value_fn=None, is_int_feature=True):
40
    self.name = name
41
    self.shape = shape
42
    self.default_value_fn = default_value_fn
43
    self.is_int_feature = is_int_feature
44

45
  def get_parsing_spec(self):
46
    return tf.FixedLenFeature(
47
        self.shape, tf.int64 if self.is_int_feature else tf.float32)
48

49
  def get_default_value(self):
50
    if self.default_value_fn:
51
      return self.default_value_fn(self.shape)
52
    else:
53
      return np.zeros(
54
          self.shape, np.int64 if self.is_int_feature else np.float32)
55