jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def check_ctc_training(self, config, input_values, *args):
        config.ctc_zero_infinity = True
        model = UniSpeechSatForCTC(config=config)
        model.to(torch_device)
        model.train()

        # freeze feature encoder
        model.freeze_feature_encoder()

        input_values = input_values[:3]

        input_lengths = [input_values.shape[-1] // i for i in [4, 2, 1]]
        max_length_labels = model._get_feat_extract_output_lengths(torch.tensor(input_lengths))
        labels = ids_tensor((input_values.shape[0], max(max_length_labels) - 2), model.config.vocab_size)

        # pad input
        for i in range(len(input_lengths)):
            input_values[i, input_lengths[i] :] = 0.0

            if max_length_labels[i] < labels.shape[-1]:
                # it's important that we make sure that target lengths are at least
                # one shorter than logit lengths to prevent -inf
                labels[i, max_length_labels[i] - 1 :] = -100

        loss = model(input_values, labels=labels).loss
        self.parent.assertFalse(torch.isinf(loss).item())

        loss.backward()
//  [python]
def convert_movilevit_checkpoint(model_name, checkpoint_path, pytorch_dump_folder_path, push_to_hub=False):
    """
    Copy/paste/tweak model's weights to our MobileNetV1 structure.
    """
    config = get_mobilenet_v1_config(model_name)

    # Load 🤗 model
    model = MobileNetV1ForImageClassification(config).eval()

    # Load weights from TensorFlow checkpoint
    load_tf_weights_in_mobilenet_v1(model, config, checkpoint_path)

    # Check outputs on an image, prepared by MobileNetV1FeatureExtractor
    feature_extractor = MobileNetV1FeatureExtractor(
        crop_size={"width": config.image_size, "height": config.image_size},
        size={"shortest_edge": config.image_size + 32},
    )
    encoding = feature_extractor(images=prepare_img(), return_tensors="pt")
    outputs = model(**encoding)
    logits = outputs.logits

    assert logits.shape == (1, 1001)

    if model_name == "mobilenet_v1_1.0_224":
        expected_logits = torch.tensor([-4.1739, -1.1233, 3.1205])
    elif model_name == "mobilenet_v1_0.75_192":
        expected_logits = torch.tensor([-3.9440, -2.3141, -0.3333])
    else:
        expected_logits = None

    if expected_logits is not None:
        assert torch.allclose(logits[0, :3], expected_logits, atol=1e-4)

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model {model_name} to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)
    print(f"Saving feature extractor to {pytorch_dump_folder_path}")
    feature_extractor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        print("Pushing to the hub...")
        repo_id = "google/" + model_name
        feature_extractor.push_to_hub(repo_id)
        model.push_to_hub(repo_id)
//  [python]
def convert_movilevit_checkpoint(model_name, checkpoint_path, pytorch_dump_folder_path, push_to_hub=False):
    """
    Copy/paste/tweak model's weights to our MobileNetV1 structure.
    """
    config = get_mobilenet_v1_config(model_name)

    # Load 🤗 model
    model = MobileNetV1ForImageClassification(config).eval()

    # Load weights from TensorFlow checkpoint
    load_tf_weights_in_mobilenet_v1(model, config, checkpoint_path)

    # Check outputs on an image, prepared by MobileNetV1ImageProcessor
    image_processor = MobileNetV1ImageProcessor(
        crop_size={"width": config.image_size, "height": config.image_size},
        size={"shortest_edge": config.image_size + 32},
    )
    encoding = image_processor(images=prepare_img(), return_tensors="pt")
    outputs = model(**encoding)
    logits = outputs.logits

    assert logits.shape == (1, 1001)

    if model_name == "mobilenet_v1_1.0_224":
        expected_logits = torch.tensor([-4.1739, -1.1233, 3.1205])
    elif model_name == "mobilenet_v1_0.75_192":
        expected_logits = torch.tensor([-3.9440, -2.3141, -0.3333])
    else:
        expected_logits = None

    if expected_logits is not None:
        assert torch.allclose(logits[0, :3], expected_logits, atol=1e-4)

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model {model_name} to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)
    print(f"Saving feature extractor to {pytorch_dump_folder_path}")
    image_processor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        print("Pushing to the hub...")
        repo_id = "google/" + model_name
        image_processor.push_to_hub(repo_id)
        model.push_to_hub(repo_id)
//  [python]
def convert_movilevit_checkpoint(
    model_name, checkpoint_path, pytorch_dump_folder_path, push_to_hub=False
):
    """
    Copy/paste/tweak model's weights to our MobileNetV1 structure.
    """
    config = get_mobilenet_v1_config(model_name)

    # Load 🤗 model
    model = MobileNetV1ForImageClassification(config).eval()

    # Load weights from TensorFlow checkpoint
    load_tf_weights_in_mobilenet_v1(model, config, checkpoint_path)

    # Check outputs on an image, prepared by MobileNetV1FeatureExtractor
    feature_extractor = MobileNetV1FeatureExtractor(
        crop_size={"width": config.image_size, "height": config.image_size},
        size={"shortest_edge": config.image_size + 32},
    )
    encoding = feature_extractor(images=prepare_img(), return_tensors="pt")
    outputs = model(**encoding)
    logits = outputs.logits

    assert logits.shape == (1, 1001)

    if model_name == "mobilenet_v1_1.0_224":
        expected_logits = torch.tensor([-4.1739, -1.1233, 3.1205])
    elif model_name == "mobilenet_v1_0.75_192":
        expected_logits = torch.tensor([-3.9440, -2.3141, -0.3333])
    else:
        expected_logits = None

    if expected_logits is not None:
        assert torch.allclose(logits[0, :3], expected_logits, atol=1e-4)

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model {model_name} to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)
    print(f"Saving feature extractor to {pytorch_dump_folder_path}")
    feature_extractor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        print("Pushing to the hub...")
        repo_id = "google/" + model_name
        feature_extractor.push_to_hub(repo_id)
        model.push_to_hub(repo_id)
//  [python]
def convert_movilevit_checkpoint(
    model_name, checkpoint_path, pytorch_dump_folder_path, push_to_hub=False
):
    """
    Copy/paste/tweak model's weights to our MobileNetV1 structure.
    """
    config = get_mobilenet_v1_config(model_name)

    # Load 🤗 model
    model = MobileNetV1ForImageClassification(config).eval()

    # Load weights from TensorFlow checkpoint
    load_tf_weights_in_mobilenet_v1(model, config, checkpoint_path)

    # Check outputs on an image, prepared by MobileNetV1ImageProcessor
    feature_extractor = MobileNetV1ImageProcessor(
        crop_size={"width": config.image_size, "height": config.image_size},
        size={"shortest_edge": config.image_size + 32},
    )
    encoding = feature_extractor(images=prepare_img(), return_tensors="pt")
    outputs = model(**encoding)
    logits = outputs.logits

    assert logits.shape == (1, 1001)

    if model_name == "mobilenet_v1_1.0_224":
        expected_logits = torch.tensor([-4.1739, -1.1233, 3.1205])
    elif model_name == "mobilenet_v1_0.75_192":
        expected_logits = torch.tensor([-3.9440, -2.3141, -0.3333])
    else:
        expected_logits = None

    if expected_logits is not None:
        assert torch.allclose(logits[0, :3], expected_logits, atol=1e-4)

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model {model_name} to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)
    print(f"Saving feature extractor to {pytorch_dump_folder_path}")
    feature_extractor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        print("Pushing to the hub...")
        repo_id = "google/" + model_name
        feature_extractor.push_to_hub(repo_id)
        model.push_to_hub(repo_id)
//  [python]
def convert_movilevit_checkpoint(model_name, checkpoint_path, pytorch_dump_folder_path, push_to_hub=False):
    """
    Copy/paste/tweak model's weights to our MobileNetV1 structure.
    """
    config = get_mobilenet_v1_config(model_name)

    # Load 🤗 model
    model = MobileNetV1ForImageClassification(config).eval()

    # Load weights from TensorFlow checkpoint
    load_tf_weights_in_mobilenet_v1(model, config, checkpoint_path)

    # Check outputs on an image, prepared by MobileNetV1FeatureExtractor
    image_processor = MobileNetV1FeatureExtractor(
        crop_size={"width": config.image_size, "height": config.image_size},
        size={"shortest_edge": config.image_size + 32},
    )
    encoding = image_processor(images=prepare_img(), return_tensors="pt")
    outputs = model(**encoding)
    logits = outputs.logits

    assert logits.shape == (1, 1001)

    if model_name == "mobilenet_v1_1.0_224":
        expected_logits = torch.tensor([-4.1739, -1.1233, 3.1205])
    elif model_name == "mobilenet_v1_0.75_192":
        expected_logits = torch.tensor([-3.9440, -2.3141, -0.3333])
    else:
        expected_logits = None

    if expected_logits is not None:
        assert torch.allclose(logits[0, :3], expected_logits, atol=1e-4)

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model {model_name} to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)
    print(f"Saving image processor to {pytorch_dump_folder_path}")
    image_processor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        print("Pushing to the hub...")
        repo_id = "google/" + model_name
        image_processor.push_to_hub(repo_id)
        model.push_to_hub(repo_id)
//  [python]
def convert_movilevit_checkpoint(model_name, checkpoint_path, pytorch_dump_folder_path, push_to_hub=False):
    """
    Copy/paste/tweak model's weights to our MobileNetV1 structure.
    """
    config = get_mobilenet_v1_config(model_name)

    # Load 🤗 model
    model = MobileNetV1ForImageClassification(config).eval()

    # Load weights from TensorFlow checkpoint
    load_tf_weights_in_mobilenet_v1(model, config, checkpoint_path)

    # Check outputs on an image, prepared by MobileNetV1ImageProcessor
    image_processor = MobileNetV1ImageProcessor(
        crop_size={"width": config.image_size, "height": config.image_size},
        size={"shortest_edge": config.image_size + 32},
    )
    encoding = image_processor(images=prepare_img(), return_tensors="pt")
    outputs = model(**encoding)
    logits = outputs.logits

    assert logits.shape == (1, 1001)

    if model_name == "mobilenet_v1_1.0_224":
        expected_logits = torch.tensor([-4.1739, -1.1233, 3.1205])
    elif model_name == "mobilenet_v1_0.75_192":
        expected_logits = torch.tensor([-3.9440, -2.3141, -0.3333])
    else:
        expected_logits = None

    if expected_logits is not None:
        assert torch.allclose(logits[0, :3], expected_logits, atol=1e-4)

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model {model_name} to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)
    print(f"Saving image processor to {pytorch_dump_folder_path}")
    image_processor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        print("Pushing to the hub...")
        repo_id = "google/" + model_name
        image_processor.push_to_hub(repo_id)
        model.push_to_hub(repo_id)
//  [python]
def convert_mobilevitv2_checkpoint(task_name, checkpoint_path, orig_config_path, pytorch_dump_folder_path):
    """
    Copy/paste/tweak model's weights to our MobileViTV2 structure.
    """
    config = get_mobilevitv2_config(task_name, orig_config_path)

    # load original state_dict
    checkpoint = torch.load(checkpoint_path, map_location="cpu")

    # load huggingface model
    if task_name.startswith("ade20k_") or task_name.startswith("voc_"):
        model = MobileViTV2ForSemanticSegmentation(config).eval()
        base_model = False
    else:
        model = MobileViTV2ForImageClassification(config).eval()
        base_model = False

    # remove and rename some keys of load the original model
    state_dict = checkpoint
    remove_unused_keys(state_dict)
    rename_keys = create_rename_keys(state_dict, base_model=base_model)
    for rename_key_src, rename_key_dest in rename_keys:
        rename_key(state_dict, rename_key_src, rename_key_dest)

    # load modified state_dict
    model.load_state_dict(state_dict)

    # Check outputs on an image, prepared by MobileViTFeatureExtractor
    image_processor = MobileViTFeatureExtractor(crop_size=config.image_size, size=config.image_size + 32)
    encoding = image_processor(images=prepare_img(), return_tensors="pt")
    outputs = model(**encoding)

    # verify classification model
    if task_name.startswith("imagenet"):
        logits = outputs.logits
        predicted_class_idx = logits.argmax(-1).item()
        print("Predicted class:", model.config.id2label[predicted_class_idx])
        if task_name.startswith("imagenet1k_256") and config.width_multiplier == 1.0:
            # expected_logits for base variant
            expected_logits = torch.tensor([-1.6336e00, -7.3204e-02, -5.1883e-01])
            assert torch.allclose(logits[0, :3], expected_logits, atol=1e-4)

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model {task_name} to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)
    print(f"Saving image processor to {pytorch_dump_folder_path}")
    image_processor.save_pretrained(pytorch_dump_folder_path)
//  [python]
def convert_mobilevitv2_checkpoint(task_name, checkpoint_path, orig_config_path, pytorch_dump_folder_path):
    """
    Copy/paste/tweak model's weights to our MobileViTV2 structure.
    """
    config = get_mobilevitv2_config(task_name, orig_config_path)

    # load original state_dict
    checkpoint = torch.load(checkpoint_path, map_location="cpu")

    # load huggingface model
    if task_name.startswith("ade20k_") or task_name.startswith("voc_"):
        model = MobileViTV2ForSemanticSegmentation(config).eval()
        base_model = False
    else:
        model = MobileViTV2ForImageClassification(config).eval()
        base_model = False

    # remove and rename some keys of load the original model
    state_dict = checkpoint
    remove_unused_keys(state_dict)
    rename_keys = create_rename_keys(state_dict, base_model=base_model)
    for rename_key_src, rename_key_dest in rename_keys:
        rename_key(state_dict, rename_key_src, rename_key_dest)

    # load modified state_dict
    model.load_state_dict(state_dict)

    # Check outputs on an image, prepared by MobileViTImageProcessor
    image_processor = MobileViTImageProcessor(crop_size=config.image_size, size=config.image_size + 32)
    encoding = image_processor(images=prepare_img(), return_tensors="pt")
    outputs = model(**encoding)

    # verify classification model
    if task_name.startswith("imagenet"):
        logits = outputs.logits
        predicted_class_idx = logits.argmax(-1).item()
        print("Predicted class:", model.config.id2label[predicted_class_idx])
        if task_name.startswith("imagenet1k_256") and config.width_multiplier == 1.0:
            # expected_logits for base variant
            expected_logits = torch.tensor([-1.6336e00, -7.3204e-02, -5.1883e-01])
            assert torch.allclose(logits[0, :3], expected_logits, atol=1e-4)

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model {task_name} to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)
    print(f"Saving image processor to {pytorch_dump_folder_path}")
    image_processor.save_pretrained(pytorch_dump_folder_path)
//  [python]
def convert_mobilevitv2_checkpoint(task_name, checkpoint_path, orig_config_path, pytorch_dump_folder_path):
    """
    Copy/paste/tweak model's weights to our MobileViTV2 structure.
    """
    config = get_mobilevitv2_config(task_name, orig_config_path)

    # load original state_dict
    checkpoint = torch.load(checkpoint_path, map_location="cpu")

    # load huggingface model
    if task_name.startswith("ade20k_") or task_name.startswith("voc_"):
        model = MobileViTV2ForSemanticSegmentation(config).eval()
        base_model = False
    else:
        model = MobileViTV2ForImageClassification(config).eval()
        base_model = False

    # remove and rename some keys of load the original model
    state_dict = checkpoint
    remove_unused_keys(state_dict)
    rename_keys = create_rename_keys(state_dict, base_model=base_model)
    for rename_key_src, rename_key_dest in rename_keys:
        rename_key(state_dict, rename_key_src, rename_key_dest)

    # load modified state_dict
    model.load_state_dict(state_dict)

    # Check outputs on an image, prepared by MobileViTFeatureExtractor
    feature_extractor = MobileViTFeatureExtractor(crop_size=config.image_size, size=config.image_size + 32)
    encoding = feature_extractor(images=prepare_img(), return_tensors="pt")
    outputs = model(**encoding)

    # verify classification model
    if task_name.startswith("imagenet"):
        logits = outputs.logits
        predicted_class_idx = logits.argmax(-1).item()
        print("Predicted class:", model.config.id2label[predicted_class_idx])
        if task_name.startswith("imagenet1k_256") and config.width_multiplier == 1.0:
            # expected_logits for base variant
            expected_logits = torch.tensor([-1.6336e00, -7.3204e-02, -5.1883e-01])
            assert torch.allclose(logits[0, :3], expected_logits, atol=1e-4)

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model {task_name} to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)
    print(f"Saving feature extractor to {pytorch_dump_folder_path}")
    feature_extractor.save_pretrained(pytorch_dump_folder_path)
//  [python]
def convert_mobilevitv2_checkpoint(task_name, checkpoint_path, orig_config_path, pytorch_dump_folder_path):
    """
    Copy/paste/tweak model's weights to our MobileViTV2 structure.
    """
    config = get_mobilevitv2_config(task_name, orig_config_path)

    # load original state_dict
    checkpoint = torch.load(checkpoint_path, map_location="cpu")

    # load huggingface model
    if task_name.startswith("ade20k_") or task_name.startswith("voc_"):
        model = MobileViTV2ForSemanticSegmentation(config).eval()
        base_model = False
    else:
        model = MobileViTV2ForImageClassification(config).eval()
        base_model = False

    # remove and rename some keys of load the original model
    state_dict = checkpoint
    remove_unused_keys(state_dict)
    rename_keys = create_rename_keys(state_dict, base_model=base_model)
    for rename_key_src, rename_key_dest in rename_keys:
        rename_key(state_dict, rename_key_src, rename_key_dest)

    # load modified state_dict
    model.load_state_dict(state_dict)

    # Check outputs on an image, prepared by MobileViTImageProcessor
    feature_extractor = MobileViTImageProcessor(crop_size=config.image_size, size=config.image_size + 32)
    encoding = feature_extractor(images=prepare_img(), return_tensors="pt")
    outputs = model(**encoding)

    # verify classification model
    if task_name.startswith("imagenet"):
        logits = outputs.logits
        predicted_class_idx = logits.argmax(-1).item()
        print("Predicted class:", model.config.id2label[predicted_class_idx])
        if task_name.startswith("imagenet1k_256") and config.width_multiplier == 1.0:
            # expected_logits for base variant
            expected_logits = torch.tensor([-1.6336e00, -7.3204e-02, -5.1883e-01])
            assert torch.allclose(logits[0, :3], expected_logits, atol=1e-4)

    Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
    print(f"Saving model {task_name} to {pytorch_dump_folder_path}")
    model.save_pretrained(pytorch_dump_folder_path)
    print(f"Saving feature extractor to {pytorch_dump_folder_path}")
    feature_extractor.save_pretrained(pytorch_dump_folder_path)
//  [python]
def test_inference_image_classification_head(self):
        feature_extractor = PoolFormerFeatureExtractor()
        model = PoolFormerForImageClassification.from_pretrained("sail/poolformer_s12").to(torch_device)

        inputs = feature_extractor(images=prepare_img(), return_tensors="pt").to(torch_device)

        # forward pass
        with torch.no_grad():
            outputs = model(**inputs)

        # verify the logits
        expected_shape = torch.Size((1, 1000))
        self.assertEqual(outputs.logits.shape, expected_shape)

        expected_slice = torch.tensor([-0.6113, 0.1685, -0.0492]).to(torch_device)
        self.assertTrue(torch.allclose(outputs.logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_head(self):
        feature_extractor = PoolFormerImageProcessor()
        model = PoolFormerForImageClassification.from_pretrained("sail/poolformer_s12").to(torch_device)

        inputs = feature_extractor(images=prepare_img(), return_tensors="pt").to(torch_device)

        # forward pass
        with torch.no_grad():
            outputs = model(**inputs)

        # verify the logits
        expected_shape = torch.Size((1, 1000))
        self.assertEqual(outputs.logits.shape, expected_shape)

        expected_slice = torch.tensor([-0.6113, 0.1685, -0.0492]).to(torch_device)
        self.assertTrue(torch.allclose(outputs.logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_head(self):
        image_processor = PoolFormerFeatureExtractor()
        model = PoolFormerForImageClassification.from_pretrained("sail/poolformer_s12").to(torch_device)

        inputs = image_processor(images=prepare_img(), return_tensors="pt").to(torch_device)

        # forward pass
        with torch.no_grad():
            outputs = model(**inputs)

        # verify the logits
        expected_shape = torch.Size((1, 1000))
        self.assertEqual(outputs.logits.shape, expected_shape)

        expected_slice = torch.tensor([-0.6113, 0.1685, -0.0492]).to(torch_device)
        self.assertTrue(torch.allclose(outputs.logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_head(self):
        image_processor = PoolFormerImageProcessor()
        model = PoolFormerForImageClassification.from_pretrained("sail/poolformer_s12").to(torch_device)

        inputs = image_processor(images=prepare_img(), return_tensors="pt").to(torch_device)

        # forward pass
        with torch.no_grad():
            outputs = model(**inputs)

        # verify the logits
        expected_shape = torch.Size((1, 1000))
        self.assertEqual(outputs.logits.shape, expected_shape)

        expected_slice = torch.tensor([-0.6113, 0.1685, -0.0492]).to(torch_device)
        self.assertTrue(torch.allclose(outputs.logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification(self):

        feature_extractor = PerceiverFeatureExtractor()
        model = PerceiverForImageClassificationLearned.from_pretrained("deepmind/vision-perceiver-learned")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = feature_extractor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1653, -0.1993, -0.7521], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification(self):

        feature_extractor = PerceiverImageProcessor()
        model = PerceiverForImageClassificationLearned.from_pretrained("deepmind/vision-perceiver-learned")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = feature_extractor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1653, -0.1993, -0.7521], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_fourier(self):

        feature_extractor = PerceiverFeatureExtractor()
        model = PerceiverForImageClassificationFourier.from_pretrained("deepmind/vision-perceiver-fourier")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = feature_extractor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1295, -0.2832, 0.3226], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_fourier(self):

        feature_extractor = PerceiverImageProcessor()
        model = PerceiverForImageClassificationFourier.from_pretrained("deepmind/vision-perceiver-fourier")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = feature_extractor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1295, -0.2832, 0.3226], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_conv(self):

        feature_extractor = PerceiverFeatureExtractor()
        model = PerceiverForImageClassificationConvProcessing.from_pretrained("deepmind/vision-perceiver-conv")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = feature_extractor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1186, 0.0554, 0.0897], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_conv(self):

        feature_extractor = PerceiverImageProcessor()
        model = PerceiverForImageClassificationConvProcessing.from_pretrained("deepmind/vision-perceiver-conv")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = feature_extractor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1186, 0.0554, 0.0897], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification(self):
        feature_extractor = PerceiverFeatureExtractor()
        model = PerceiverForImageClassificationLearned.from_pretrained("deepmind/vision-perceiver-learned")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = feature_extractor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1652, -0.1992, -0.7520], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification(self):
        feature_extractor = PerceiverImageProcessor()
        model = PerceiverForImageClassificationLearned.from_pretrained("deepmind/vision-perceiver-learned")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = feature_extractor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1652, -0.1992, -0.7520], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_fourier(self):
        feature_extractor = PerceiverFeatureExtractor()
        model = PerceiverForImageClassificationFourier.from_pretrained("deepmind/vision-perceiver-fourier")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = feature_extractor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1295, -0.2832, 0.3226], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_fourier(self):
        feature_extractor = PerceiverImageProcessor()
        model = PerceiverForImageClassificationFourier.from_pretrained("deepmind/vision-perceiver-fourier")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = feature_extractor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1295, -0.2832, 0.3226], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_conv(self):
        feature_extractor = PerceiverFeatureExtractor()
        model = PerceiverForImageClassificationConvProcessing.from_pretrained("deepmind/vision-perceiver-conv")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = feature_extractor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1186, 0.0554, 0.0897], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_conv(self):
        feature_extractor = PerceiverImageProcessor()
        model = PerceiverForImageClassificationConvProcessing.from_pretrained("deepmind/vision-perceiver-conv")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = feature_extractor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1186, 0.0554, 0.0897], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def __init__(
        self, 
        hf_dataset: datasets.arrow_dataset.Dataset
        ) -> None:
        super().__init__()
        self.dataset = hf_dataset
        self.feature_extractor = PerceiverFeatureExtractor()
//  [python]
def __init__(
        self, 
        hf_dataset: datasets.arrow_dataset.Dataset
        ) -> None:
        super().__init__()
        self.dataset = hf_dataset
        self.feature_extractor = PerceiverImageProcessor()
//  [python]
def test_inference_image_classification(self):
        image_processor = PerceiverFeatureExtractor()
        model = PerceiverForImageClassificationLearned.from_pretrained("deepmind/vision-perceiver-learned")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = image_processor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1652, -0.1992, -0.7520], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification(self):
        image_processor = PerceiverImageProcessor()
        model = PerceiverForImageClassificationLearned.from_pretrained("deepmind/vision-perceiver-learned")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = image_processor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1652, -0.1992, -0.7520], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_fourier(self):
        image_processor = PerceiverFeatureExtractor()
        model = PerceiverForImageClassificationFourier.from_pretrained("deepmind/vision-perceiver-fourier")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = image_processor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1295, -0.2832, 0.3226], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_fourier(self):
        image_processor = PerceiverImageProcessor()
        model = PerceiverForImageClassificationFourier.from_pretrained("deepmind/vision-perceiver-fourier")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = image_processor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1295, -0.2832, 0.3226], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_conv(self):
        image_processor = PerceiverFeatureExtractor()
        model = PerceiverForImageClassificationConvProcessing.from_pretrained("deepmind/vision-perceiver-conv")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = image_processor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1186, 0.0554, 0.0897], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_inference_image_classification_conv(self):
        image_processor = PerceiverImageProcessor()
        model = PerceiverForImageClassificationConvProcessing.from_pretrained("deepmind/vision-perceiver-conv")
        model.to(torch_device)

        # prepare inputs
        image = prepare_img()
        inputs = image_processor(image, return_tensors="pt").pixel_values.to(torch_device)
        input_mask = None

        # forward pass
        with torch.no_grad():
            outputs = model(inputs=inputs, attention_mask=input_mask)
        logits = outputs.logits

        # verify logits
        expected_shape = torch.Size((1, model.config.num_labels))
        self.assertEqual(logits.shape, expected_shape)

        expected_slice = torch.tensor([-1.1186, 0.0554, 0.0897], device=torch_device)

        self.assertTrue(torch.allclose(logits[0, :3], expected_slice, atol=1e-4))
//  [python]
def test_multiple_init():
    # Test that multiple inits does not much worse than a single one
    X = rng.randn(30, 5)
    X[:10] += 2
    g = mixture.GMM(n_components=2, covariance_type='spherical',
                    random_state=rng, min_covar=1e-7, n_iter=5)
    with ignore_warnings(category=DeprecationWarning):
        train1 = g.fit(X).score(X).sum()
        g.n_init = 5
        train2 = g.fit(X).score(X).sum()
    assert_true(train2 >= train1 - 1.e-2)
//  [python]
def test_multiple_init():
    # Test that multiple inits does not much worse than a single one
    X = rng.randn(30, 5)
    X[:10] += 2
    g = mixture.GaussianMixture(n_components=2, covariance_type='spherical',
                    random_state=rng, min_covar=1e-7, n_init=5, max_iter=5)
    with ignore_warnings(category=DeprecationWarning):
        train1 = g.fit(X).score(X).sum()
        g.n_init = 5
        train2 = g.fit(X).score(X).sum()
    assert_true(train2 >= train1 - 1.e-2)
//  [python]
def test_n_parameters():
    n_samples, n_dim, n_components = 7, 5, 2
    X = rng.randn(n_samples, n_dim)
    n_params = {'spherical': 13, 'diag': 21, 'tied': 26, 'full': 41}
    for cv_type in ['full', 'tied', 'diag', 'spherical']:
        with ignore_warnings(category=DeprecationWarning):
            g = mixture.GMM(n_components=n_components, covariance_type=cv_type,
                            random_state=rng, min_covar=1e-7, n_iter=1)
            g.fit(X)
            assert_true(g._n_parameters() == n_params[cv_type])
//  [python]
def test_n_parameters():
    n_samples, n_dim, n_components = 7, 5, 2
    X = rng.randn(n_samples, n_dim)
    n_params = {'spherical': 13, 'diag': 21, 'tied': 26, 'full': 41}
    for cv_type in ['full', 'tied', 'diag', 'spherical']:
        with ignore_warnings(category=DeprecationWarning):
            g = mixture.GaussianMixture(n_components=n_components, covariance_type=cv_type,
                            random_state=rng, min_covar=1e-7, max_iter=1)
            g.fit(X)
            assert_true(g._n_parameters() == n_params[cv_type])
//  [python]
def test_1d_1component():
    # Test all of the covariance_types return the same BIC score for
    # 1-dimensional, 1 component fits.
    n_samples, n_dim, n_components = 100, 1, 1
    X = rng.randn(n_samples, n_dim)
    g_full = mixture.GMM(n_components=n_components, covariance_type='full',
                         random_state=rng, min_covar=1e-7, n_iter=1)
    with ignore_warnings(category=DeprecationWarning):
        g_full.fit(X)
        g_full_bic = g_full.bic(X)
        for cv_type in ['tied', 'diag', 'spherical']:
            g = mixture.GMM(n_components=n_components, covariance_type=cv_type,
                            random_state=rng, min_covar=1e-7, n_iter=1)
            g.fit(X)
            assert_array_almost_equal(g.bic(X), g_full_bic)
//  [python]
def test_1d_1component():
    # Test all of the covariance_types return the same BIC score for
    # 1-dimensional, 1 component fits.
    n_samples, n_dim, n_components = 100, 1, 1
    X = rng.randn(n_samples, n_dim)
    g_full = mixture.GaussianMixture(n_components=n_components, covariance_type='full',
                         random_state=rng, min_covar=1e-7, max_iter=1)
    g_full.fit(X)
    g_full_bic = g_full.bic(X)
    for cv_type in ['tied', 'diag', 'spherical']:
        g = mixture.GaussianMixture(n_components=n_components, covariance_type=cv_type,
                            random_state=rng, min_covar=1e-7, max_iter=1)
        g.fit(X)
        assert_array_almost_equal(g.bic(X), g_full_bic)
//  [python]
def test_GMM_attributes():
    n_components, n_features = 10, 4
    covariance_type = 'diag'
    g = mixture.GMM(n_components, covariance_type, random_state=rng)
    weights = rng.rand(n_components)
    weights = weights / weights.sum()
    means = rng.randint(-20, 20, (n_components, n_features))

    assert_true(g.n_components == n_components)
    assert_true(g.covariance_type == covariance_type)

    g.weights_ = weights
    assert_array_almost_equal(g.weights_, weights)
    g.means_ = means
    assert_array_almost_equal(g.means_, means)

    covars = (0.1 + 2 * rng.rand(n_components, n_features)) ** 2
    g.covars_ = covars
    assert_array_almost_equal(g.covars_, covars)
    assert_raises(ValueError, g._set_covars, [])
    assert_raises(ValueError, g._set_covars,
                  np.zeros((n_components - 2, n_features)))

    assert_raises(ValueError, mixture.GMM, n_components=20,
                  covariance_type='badcovariance_type')
//  [python]
def test_GMM_attributes():
    n_components, n_features = 10, 4
    covariance_type = 'diag'
    g = mixture.GaussianMixture(n_components, covariance_type, random_state=rng)
    weights = rng.rand(n_components)
    weights = weights / weights.sum()
    means = rng.randint(-20, 20, (n_components, n_features))

    assert_true(g.n_components == n_components)
    assert_true(g.covariance_type == covariance_type)

    g.weights_ = weights
    assert_array_almost_equal(g.weights_, weights)
    g.means_ = means
    assert_array_almost_equal(g.means_, means)

    covars = (0.1 + 2 * rng.rand(n_components, n_features)) ** 2
    g.covariances_ = covars
    assert_array_almost_equal(g.covariances_, covars)
    assert_raises(ValueError, g._set_covars, [])
    assert_raises(ValueError, g._set_covars,
                  np.zeros((n_components - 2, n_features)))

    assert_raises(ValueError, mixture.GaussianMixture, n_components=20,
                  covariance_type='badcovariance_type')
//  [python]
def test_fit_predict():
    """
    test that gmm.fit_predict is equivalent to gmm.fit + gmm.predict
    """
    lrng = np.random.RandomState(101)

    n_samples, n_dim, n_comps = 100, 2, 2
    mu = np.array([[8, 8]])
    component_0 = lrng.randn(n_samples, n_dim)
    component_1 = lrng.randn(n_samples, n_dim) + mu
    X = np.vstack((component_0, component_1))

    for m_constructor in (mixture.GaussianMixture, mixture.BayesianGaussianMixture, mixture.BayesianGaussianMixture):
        model = m_constructor(n_components=n_comps, covariance_type='full',
                              min_covar=1e-7, max_iter=5,
                              random_state=np.random.RandomState(0))
        assert_fit_predict_correct(model, X)

    model = mixture.GaussianMixture(n_components=n_comps, max_iter=0)
    z = model.fit_predict(X)
    assert np.all(z == 0), "Quick Initialization Failed!"
//  [python]
def test_verbose_first_level():
    # Create sample data
    X = rng.randn(30, 5)
    X[:10] += 2
    g = mixture.GMM(n_components=2, n_init=2, verbose=1)

    old_stdout = sys.stdout
    sys.stdout = StringIO()
    try:
        g.fit(X)
    finally:
        sys.stdout = old_stdout
//  [python]
def test_verbose_first_level():
    # Create sample data
    X = rng.randn(30, 5)
    X[:10] += 2
    g = mixture.GaussianMixture(n_components=2, n_init=2, verbose=1)

    old_stdout = sys.stdout
    sys.stdout = StringIO()
    try:
        g.fit(X)
    finally:
        sys.stdout = old_stdout