import six


def call_module(module, one_hots, row_lengths, signature):
  """Call a tf_hub.Module using the standard blundell signature.

  This expects that `module` has a signature named `signature` which conforms to
      ('sequence',
       'sequence_length') -> output

  To use an existing SavedModel
  file you may want to create a module_spec with
  `tensorflow_hub.saved_model_module.create_module_spec_from_saved_model`.

  Args:
    module: a tf_hub.Module to call.
    one_hots: a rank 3 tensor with one-hot encoded sequences of residues.
    row_lengths: a rank 1 tensor with sequence lengths.
    signature: the graph signature to validate and call.

  Returns:
    The output tensor of `module`.
  """
  if signature not in module.get_signature_names():
    raise ValueError('signature not in ' +
                     six.ensure_str(str(module.get_signature_names())) +
                     '. Was ' + six.ensure_str(signature) + '.')
  inputs = module.get_input_info_dict(signature=signature)
  expected_inputs = [
      'sequence',
      'sequence_length',
  ]
  if set(inputs.keys()) != set(expected_inputs):
    raise ValueError(
        'The signature_def does not have the expected inputs. Please '
        'reconfigure your saved model to only export signatures '
        'with sequence and length inputs. (Inputs were %s, expected %s)' %
        (str(inputs), str(expected_inputs)))

  outputs = module.get_output_info_dict(signature=signature)
  if len(outputs) > 1:
    raise ValueError('The signature_def given has more than one output. Please '
                     'reconfigure your saved model to only export signatures '
                     'with one output. (Outputs were %s)' % str(outputs))

  return list(
      module({
          'sequence': one_hots,
          'sequence_length': row_lengths,
      },
             signature=signature,
             as_dict=True).values())[0]