<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 See code at https
Step1:  Use model
Step2:  tfds zero-shot evaluation
<ASSISTANT_TASK:>
Python Code:
# Installs the vit_jax package from Github.
!pip install -q git+https://github.com/google-research/vision_transformer
import jax
import jax.numpy as jnp
from matplotlib import pyplot as plt
import numpy as np
import pandas as pd
import tensorflow as tf
import tensorflow_datasets as tfds
import tqdm
from vit_jax import models
# Currently available LiT models
[name for name in models.model_configs.MODEL_CONFIGS if name.startswith('LiT')]
model_name = 'LiT-B16B'
lit_model = models.get_model(model_name)
# Loading the variables from cloud can take a while the first time...
lit_variables = lit_model.load_variables()
# Creating tokens from freeform text (see next section).
tokenizer = lit_model.get_tokenizer()
# Resizing images & converting value range to -1..1 (see next section).
image_preprocessing = lit_model.get_image_preprocessing()
# Preprocessing op for use in tfds pipeline (see last section).
pp = lit_model.get_pp()
Explanation: See code at https://github.com/google-research/vision_transformer/
This Colab is about the paper
LiT: Zero-Shot Transfer with Locked-image text Tuning: https://arxiv.org/abs/2111.07991
For ViT, MLP Mixer etc see the other Colab
https://colab.research.google.com/github/google-research/vision_transformer/blob/main/vit_jax.ipynb
Load model
End of explanation
# Let's load some sample images from tfds.
# Alternatively you can also load these images from the internet / your Drive.
ds = tfds.load('imagenette', split='train')
images_list = [
    example['image'].numpy()
    for _, example in zip(range(5), ds)
]
# Note that this is a list of images with different shapes, not a four
# dimensional tensor.
[image.shape for image in images_list]
# Note that our preprocessing converts to floats ranging from -1..1 !
images = image_preprocessing(images_list)
images.shape, images.min(), images.max()
plt.figure(figsize=(15, 4))
plt.imshow(np.hstack(images) * .5 + .5)
plt.axis('off');
texts = [
    'itap of a cd player',
    'a photo of a truck',
    'gas station',
    'chainsaw',
    'a bad photo of colorful houses',
]
tokens = tokenizer(texts)
tokens.shape
# Embed both texts and images with a single model call.
# See next section for embedding images/texts separately.
zimg, ztxt, out = lit_model.apply(lit_variables, images=images, tokens=tokens)
plt.imshow(ztxt @ zimg.T)
probs = np.array(jax.nn.softmax(out['t'] * ztxt @ zimg.T, axis=1))
pd.DataFrame(probs, index=texts).style.background_gradient('Greens', vmin=0, vmax=1).format('{:.2%}')
Explanation: Use model
End of explanation
# Load dataset and create array of class names.
builder = tfds.builder('cifar100')
builder.download_and_prepare()
ds_test = builder.as_dataset('test')
info = builder.info
classnames = [
    info.features['label'].int2str(id_)
    for id_ in range(info.features['label'].num_classes)
]
classnames[:10]
# "best prompts" from CLIP paper (https://arxiv.org/abs/2103.00020)
PROMPTS = [
    'itap of a {}.',
    'a bad photo of the {}.',
    'a origami {}.',
    'a photo of the large {}.',
    'a {} in a video game.',
    'art of the {}.',
    'a photo of the small {}.',
    '{}',
]
texts = [
    prompt.format(classname)
    for classname in classnames
    for prompt in PROMPTS
]
len(texts)
# Tokenize the texts using numpy like before.
tokens = tokenizer(texts)
tokens.shape
_, ztxt, _ = lit_model.apply(lit_variables, tokens=tokens)
ztxt.shape
# `pp` from above (section "Load model") is a TensorFlow graph that can
# efficiently be added to the input pre-processing.
imgs = next(iter(ds_test.map(pp).batch(4)))['image']
# Note that `pp` would also tokenize "texts" to "tokens", if such a feature was
# present in the dataset (which is not the case for cifar).
plt.figure(figsize=(15, 4))
plt.imshow(np.hstack(imgs) * .5 + .5)
plt.axis('off');
# JIT-compile image embedding function because there are lots of images.
@jax.jit
def embed_images(variables, images):
  zimg, _, _ = lit_model.apply(variables, images=images)
  return zimg
# Compute all images embeddings & collect correct labels.
zimgs = []
labels = []
for batch in tqdm.tqdm(ds_test.map(lit_model.get_pp()).batch(500)):
  labels += list(batch['label'].numpy())
  zimg = embed_images(lit_variables, batch['image'].numpy())
  zimgs.append(np.array(zimg))
zimgs = np.concatenate(zimgs)
zimgs.shape
# Compute similarities ...
sims = zimgs @ ztxt.reshape([len(classnames), len(PROMPTS), -1]).mean(axis=1).T
sims.shape
# ... and use most similar embedding to predict label.
(sims.argmax(axis=1) == np.array(labels)).mean()
# Expected accuracy for model "LiT-B16B" : 79.19
Explanation: tfds zero-shot evaluation
End of explanation
<END_TASK>