Error running export to tf colab

[MaxTeselkin]

Hi! I tried to run you colab for exporting OWL-ViT to tensorflow, but got the following error:

TypeError: call_module() got an unexpected keyword argument 'function_list'

[guarin]

I was able to run the colab by using tensorflow 2.14 and modifying the predict_fn function.

Update first cell to:

!rm -rf *
!rm -rf .config
!rm -rf .git
!git clone https://github.com/google-research/scenic.git .
!python -m pip install -q .
!python -m pip install -r ./scenic/projects/owl_vit/requirements.txt

# Also install big_vision, which is needed for the mask head:
!mkdir /big_vision
!git clone https://github.com/google-research/big_vision.git /big_vision
!python -m pip install -r /big_vision/big_vision/requirements.txt
import sys
sys.path.append('/big_vision/')
!echo "Done."

# Use new tensorflow version
!pip install tensorflow==2.14 tensorflow-text==2.14 numpy==1.23.5

And change predict_fn to:

def predict_fn(variables, inputs):
  """Calls the model. The keys of `inputs` determine the call signature."""

  # Default signature: End-to-end prediction:
  if set(inputs.keys()) == {'images', 'tokenized_queries'}:
    return module.apply(
        variables,
        inputs=inputs['images'],
        text_queries=inputs['tokenized_queries'],
        train=False,
        true_boxes=None,
      )

  # Only images are provided: Get image embeddings:
  elif set(inputs.keys()) == {'images'}:
    return module.apply(
        variables,
        images=inputs['images'],
        train=False,
        method=module.image_embedder
    )

  # Only queries are provided: Get query embeddings:
  elif set(inputs.keys()) == {'tokenized_queries'}:
    return module.apply(
        variables,
        text_queries=inputs['tokenized_queries'],
        train=False,
        method=module.text_embedder)

  # Image features are provided: Get bounding boxes:
  elif set(inputs.keys()) == {'feature_map'}:
    shape = inputs['feature_map'].shape
    return module.apply(
        variables,
        image_features=jnp.reshape(inputs['feature_map'], (shape[0], -1, shape[-1])),
        feature_map=inputs['feature_map'],
        method=module.box_predictor)

  # Image and query embeddings are provided: Get classification scores:
  elif set(inputs.keys()) == {'feature_map', 'query_embeddings'}:
    shape = inputs['feature_map'].shape
    return module.apply(
        variables,
        image_features=jnp.reshape(inputs['feature_map'], (shape[0], -1, shape[-1])),
        query_embeddings=inputs['query_embeddings'],
        query_mask=None,
        method=module.class_predictor)

  else:
    raise ValueError(f'Unknown input signature with keys {inputs.keys()}')

[MaxTeselkin]

@guarin thanks