jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def convert_vit_msn_checkpoint(checkpoint_url, pytorch_dump_folder_path):
    config = ViTMSNConfig()
    config.num_labels = 1000

    repo_id = "datasets/huggingface/label-files"
    filename = "imagenet-1k-id2label.json"
    id2label = json.load(open(hf_hub_download(repo_id, filename), "r"))
    id2label = {int(k): v for k, v in id2label.items()}
    config.id2label = id2label
    config.label2id = {v: k for k, v in id2label.items()}

    if "s16" in checkpoint_url:
        config.hidden_size = 384
        config.intermediate_size = 1536
        config.num_attention_heads = 6
    elif "l16" in checkpoint_url:
        config.hidden_size = 1024
        config.intermediate_size = 4096
        config.num_hidden_layers = 24
        config.num_attention_heads = 16
        config.hidden_dropout_prob = 0.1
    elif "b4" in checkpoint_url:
        config.patch_size = 4
    elif "l7" in checkpoint_url:
        config.patch_size = 7
        config.hidden_size = 1024
        config.intermediate_size = 4096
        config.num_hidden_layers = 24
        config.num_attention_heads = 16
        config.hidden_dropout_prob = 0.1

    model = ViTMSNModel(config)

    state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")["target_encoder"]

    image_processor = ViTFeatureExtractor(size=config.image_size)

    remove_projection_head(state_dict)
    rename_keys = create_rename_keys(config, base_model=True)

    for src, dest in rename_keys:
        rename_key(state_dict, src, dest)
    read_in_q_k_v(state_dict, config, base_model=True)

    model.load_state_dict(state_dict)
    model.eval()

    url = "http://images.cocodataset.org/val2017/000000039769.jpg"

    image = Image.open(requests.get(url, stream=True).raw)
    image_processor = ViTFeatureExtractor(
        size=config.image_size, image_mean=IMAGENET_DEFAULT_MEAN, image_std=IMAGENET_DEFAULT_STD
    )
    inputs = image_processor(images=image, return_tensors="pt")

    # forward pass
    torch.manual_seed(2)
    outputs = model(**inputs)
    last_hidden_state = outputs.last_hidden_state

    # The following Colab Notebook was used to generate these outputs:
    # https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
    if "s16" in checkpoint_url:
        expected_slice = torch.tensor([[-1.0915, -1.4876, -1.1809]])
    elif "b16" in checkpoint_url:
        expected_slice = torch.tensor([[14.2889, -18.9045, 11.7281]])
    elif "l16" in checkpoint_url:
        expected_slice = torch.tensor([[41.5028, -22.8681, 45.6475]])
    elif "b4" in checkpoint_url:
        expected_slice = torch.tensor([[-4.3868, 5.2932, -0.4137]])
    else:
        expected_slice = torch.tensor([[-0.1792, -0.6465, 2.4263]])

    # verify logits
    assert torch.allclose(last_hidden_state[:, 0, :3], expected_slice, atol=1e-4)

    print(f"Saving model 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_vit_msn_checkpoint(checkpoint_url, pytorch_dump_folder_path):
    config = ViTMSNConfig()
    config.num_labels = 1000

    repo_id = "datasets/huggingface/label-files"
    filename = "imagenet-1k-id2label.json"
    id2label = json.load(open(hf_hub_download(repo_id, filename), "r"))
    id2label = {int(k): v for k, v in id2label.items()}
    config.id2label = id2label
    config.label2id = {v: k for k, v in id2label.items()}

    if "s16" in checkpoint_url:
        config.hidden_size = 384
        config.intermediate_size = 1536
        config.num_attention_heads = 6
    elif "l16" in checkpoint_url:
        config.hidden_size = 1024
        config.intermediate_size = 4096
        config.num_hidden_layers = 24
        config.num_attention_heads = 16
        config.hidden_dropout_prob = 0.1
    elif "b4" in checkpoint_url:
        config.patch_size = 4
    elif "l7" in checkpoint_url:
        config.patch_size = 7
        config.hidden_size = 1024
        config.intermediate_size = 4096
        config.num_hidden_layers = 24
        config.num_attention_heads = 16
        config.hidden_dropout_prob = 0.1

    model = ViTMSNModel(config)

    state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")["target_encoder"]

    image_processor = ViTImageProcessor(size=config.image_size)

    remove_projection_head(state_dict)
    rename_keys = create_rename_keys(config, base_model=True)

    for src, dest in rename_keys:
        rename_key(state_dict, src, dest)
    read_in_q_k_v(state_dict, config, base_model=True)

    model.load_state_dict(state_dict)
    model.eval()

    url = "http://images.cocodataset.org/val2017/000000039769.jpg"

    image = Image.open(requests.get(url, stream=True).raw)
    image_processor = ViTImageProcessor(
        size=config.image_size, image_mean=IMAGENET_DEFAULT_MEAN, image_std=IMAGENET_DEFAULT_STD
    )
    inputs = image_processor(images=image, return_tensors="pt")

    # forward pass
    torch.manual_seed(2)
    outputs = model(**inputs)
    last_hidden_state = outputs.last_hidden_state

    # The following Colab Notebook was used to generate these outputs:
    # https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
    if "s16" in checkpoint_url:
        expected_slice = torch.tensor([[-1.0915, -1.4876, -1.1809]])
    elif "b16" in checkpoint_url:
        expected_slice = torch.tensor([[14.2889, -18.9045, 11.7281]])
    elif "l16" in checkpoint_url:
        expected_slice = torch.tensor([[41.5028, -22.8681, 45.6475]])
    elif "b4" in checkpoint_url:
        expected_slice = torch.tensor([[-4.3868, 5.2932, -0.4137]])
    else:
        expected_slice = torch.tensor([[-0.1792, -0.6465, 2.4263]])

    # verify logits
    assert torch.allclose(last_hidden_state[:, 0, :3], expected_slice, atol=1e-4)

    print(f"Saving model 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_clipseg_checkpoint(model_name, checkpoint_path, pytorch_dump_folder_path, push_to_hub):
    config = get_clipseg_config(model_name)
    model = CLIPSegForImageSegmentation(config)
    model.eval()

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

    # remove some keys
    for key in state_dict.copy().keys():
        if key.startswith("model"):
            state_dict.pop(key, None)

    # rename some keys
    state_dict = convert_state_dict(state_dict, config)
    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)

    if missing_keys != ["clip.text_model.embeddings.position_ids", "clip.vision_model.embeddings.position_ids"]:
        raise ValueError("Missing keys that are not expected: {}".format(missing_keys))
    if unexpected_keys != ["decoder.reduce.weight", "decoder.reduce.bias"]:
        raise ValueError(f"Unexpected keys: {unexpected_keys}")

    image_processor = ViTFeatureExtractor(size=352)
    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
    processor = CLIPSegProcessor(image_processor=image_processor, tokenizer=tokenizer)

    image = prepare_img()
    text = ["a glass", "something to fill", "wood", "a jar"]

    inputs = processor(text=text, images=[image] * len(text), padding="max_length", return_tensors="pt")

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

    # verify values
    expected_conditional = torch.tensor([0.1110, -0.1882, 0.1645])
    expected_pooled_output = torch.tensor([0.2692, -0.7197, -0.1328])
    if model_name == "clipseg-rd64-refined":
        expected_masks_slice = torch.tensor(
            [[-10.0407, -9.9431, -10.2646], [-9.9751, -9.7064, -9.9586], [-9.6891, -9.5645, -9.9618]]
        )
    elif model_name == "clipseg-rd64":
        expected_masks_slice = torch.tensor(
            [[-7.2877, -7.2711, -7.2463], [-7.2652, -7.2780, -7.2520], [-7.2239, -7.2204, -7.2001]]
        )
    elif model_name == "clipseg-rd16":
        expected_masks_slice = torch.tensor(
            [[-6.3955, -6.4055, -6.4151], [-6.3911, -6.4033, -6.4100], [-6.3474, -6.3702, -6.3762]]
        )
    else:
        raise ValueError(f"Model name {model_name} not supported.")

    assert torch.allclose(outputs.logits[0, :3, :3], expected_masks_slice, atol=1e-3)
    assert torch.allclose(outputs.conditional_embeddings[0, :3], expected_conditional, atol=1e-3)
    assert torch.allclose(outputs.pooled_output[0, :3], expected_pooled_output, atol=1e-3)
    print("Looks ok!")

    if pytorch_dump_folder_path is not None:
        print(f"Saving model and processor to {pytorch_dump_folder_path}")
        model.save_pretrained(pytorch_dump_folder_path)
        processor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        print(f"Pushing model and processor for {model_name} to the hub")
        model.push_to_hub(f"CIDAS/{model_name}")
        processor.push_to_hub(f"CIDAS/{model_name}")
//  [python]
def convert_clipseg_checkpoint(model_name, checkpoint_path, pytorch_dump_folder_path, push_to_hub):
    config = get_clipseg_config(model_name)
    model = CLIPSegForImageSegmentation(config)
    model.eval()

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

    # remove some keys
    for key in state_dict.copy().keys():
        if key.startswith("model"):
            state_dict.pop(key, None)

    # rename some keys
    state_dict = convert_state_dict(state_dict, config)
    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)

    if missing_keys != ["clip.text_model.embeddings.position_ids", "clip.vision_model.embeddings.position_ids"]:
        raise ValueError("Missing keys that are not expected: {}".format(missing_keys))
    if unexpected_keys != ["decoder.reduce.weight", "decoder.reduce.bias"]:
        raise ValueError(f"Unexpected keys: {unexpected_keys}")

    image_processor = ViTImageProcessor(size=352)
    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
    processor = CLIPSegProcessor(image_processor=image_processor, tokenizer=tokenizer)

    image = prepare_img()
    text = ["a glass", "something to fill", "wood", "a jar"]

    inputs = processor(text=text, images=[image] * len(text), padding="max_length", return_tensors="pt")

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

    # verify values
    expected_conditional = torch.tensor([0.1110, -0.1882, 0.1645])
    expected_pooled_output = torch.tensor([0.2692, -0.7197, -0.1328])
    if model_name == "clipseg-rd64-refined":
        expected_masks_slice = torch.tensor(
            [[-10.0407, -9.9431, -10.2646], [-9.9751, -9.7064, -9.9586], [-9.6891, -9.5645, -9.9618]]
        )
    elif model_name == "clipseg-rd64":
        expected_masks_slice = torch.tensor(
            [[-7.2877, -7.2711, -7.2463], [-7.2652, -7.2780, -7.2520], [-7.2239, -7.2204, -7.2001]]
        )
    elif model_name == "clipseg-rd16":
        expected_masks_slice = torch.tensor(
            [[-6.3955, -6.4055, -6.4151], [-6.3911, -6.4033, -6.4100], [-6.3474, -6.3702, -6.3762]]
        )
    else:
        raise ValueError(f"Model name {model_name} not supported.")

    assert torch.allclose(outputs.logits[0, :3, :3], expected_masks_slice, atol=1e-3)
    assert torch.allclose(outputs.conditional_embeddings[0, :3], expected_conditional, atol=1e-3)
    assert torch.allclose(outputs.pooled_output[0, :3], expected_pooled_output, atol=1e-3)
    print("Looks ok!")

    if pytorch_dump_folder_path is not None:
        print(f"Saving model and processor to {pytorch_dump_folder_path}")
        model.save_pretrained(pytorch_dump_folder_path)
        processor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        print(f"Pushing model and processor for {model_name} to the hub")
        model.push_to_hub(f"CIDAS/{model_name}")
        processor.push_to_hub(f"CIDAS/{model_name}")
//  [python]
def performance_check():
    torch_model = ERNIEVIL2()
    ckpt = torch.load('ERNIE_VIL2_BASE_ViT_convert.pt', map_location='cpu')
    torch_model.load_state_dict(ckpt)
    torch_model.eval()
    image_processor = ViTFeatureExtractor(image_mean=[0.485, 0.456, 0.406],
                  image_std= [0.229, 0.224, 0.225])
    tokenizer = transformers.BertTokenizer.from_pretrained('nghuyong/ernie-3.0-base-zh') 
    
    image_path = ['src/cat.jpg','src/dog.jpg', 'src/cyberpunk.jpg']
    images = [Image.open(each).convert('RGB') for each in image_path]
    image_input = image_processor(images=images, return_tensors="pt")
    
    texts = ['一张猫的照片', '一张狗的照片',"赛博朋克边缘行者"]    
    inputs = tokenizer(text=texts, padding=True, return_tensors="pt")
    
    with torch.no_grad():
        visual_output = torch_model.visual(image_input['pixel_values'])[:, 0]
        title_features = torch_model.ernie(**inputs).pooler_output
    
    title_features = title_features / title_features.norm(p=2, dim=-1, keepdim=True)  # normalize
    visual_output = visual_output / visual_output.norm(p=2, dim=-1, keepdim=True)  # normalize
    print(title_features @ visual_output.T)
//  [python]
def performance_check():
    torch_model = ERNIEVIL2()
    ckpt = torch.load('ERNIE_VIL2_BASE_ViT_convert.pt', map_location='cpu')
    torch_model.load_state_dict(ckpt)
    torch_model.eval()
    image_processor = ViTImageProcessor(image_mean=[0.485, 0.456, 0.406],
                  image_std= [0.229, 0.224, 0.225])
    tokenizer = transformers.BertTokenizer.from_pretrained('nghuyong/ernie-3.0-base-zh') 
    
    image_path = ['src/cat.jpg','src/dog.jpg', 'src/cyberpunk.jpg']
    images = [Image.open(each).convert('RGB') for each in image_path]
    image_input = image_processor(images=images, return_tensors="pt")
    
    texts = ['一张猫的照片', '一张狗的照片',"赛博朋克边缘行者"]    
    inputs = tokenizer(text=texts, padding=True, return_tensors="pt")
    
    with torch.no_grad():
        visual_output = torch_model.visual(image_input['pixel_values'])[:, 0]
        title_features = torch_model.ernie(**inputs).pooler_output
    
    title_features = title_features / title_features.norm(p=2, dim=-1, keepdim=True)  # normalize
    visual_output = visual_output / visual_output.norm(p=2, dim=-1, keepdim=True)  # normalize
    print(title_features @ visual_output.T)
//  [python]
def logit_check(args):
    torch_model = ERNIEVIL2()
    ckpt = torch.load('ERNIE_VIL2_BASE_ViT_convert.pt', map_location='cpu')
    torch_model.load_state_dict(ckpt)
    torch_model.eval()
    ## torch_visual
    image_processor = ViTFeatureExtractor(image_mean=[0.485, 0.456, 0.406],
                  image_std= [0.229, 0.224, 0.225])
    
    image = Image.open('src/dog.jpg').convert('RGB')
    image_input = image_processor(images=[image], return_tensors="pt")
    with torch.no_grad():
        t_visual_output = torch_model.visual(image_input['pixel_values'])
    
    ## torch_text
    tokenizer = transformers.BertTokenizer.from_pretrained('nghuyong/ernie-3.0-base-zh')

    input_ids = torch.tensor([tokenizer.encode(text="welcome to ernie pytorch project", add_special_tokens=True)])
    with torch.no_grad():
        t_text_output = torch_model.ernie(input_ids)
    print("### pytorch result")
    print("visual_output: ", t_visual_output[:,0][:,:10].numpy())
    print("text_output: ", t_text_output.pooler_output[:,:10].numpy())
    
    ## paddle_visual
    place = "cpu"
    paddle.set_device(place)
    paddle_model = ERNIE_ViL2_base(args)
    paddle_model.eval()
    with paddle.no_grad():
        p_visual_output = paddle_model.visual(paddle.to_tensor(image_input['pixel_values'].numpy()))

    tokenizer = paddlenlp.transformers.AutoTokenizer.from_pretrained("ernie-3.0-base-zh")
    inputs = tokenizer("welcome to ernie pytorch project")
    inputs = {k: paddle.to_tensor([v]) for (k, v) in inputs.items()}
    
    with paddle.no_grad():
        p_text_output = paddle_model.text_model(inputs['input_ids'])
    import pdb; pdb.set_trace()
    print("### paddle result")
    print("visual_output: ", p_visual_output[:,0][:,:10].numpy())
    print("text_output: ", p_text_output[0][:,:10].numpy())
//  [python]
def logit_check(args):
    torch_model = ERNIEVIL2()
    ckpt = torch.load('ERNIE_VIL2_BASE_ViT_convert.pt', map_location='cpu')
    torch_model.load_state_dict(ckpt)
    torch_model.eval()
    ## torch_visual
    image_processor = ViTImageProcessor(image_mean=[0.485, 0.456, 0.406],
                  image_std= [0.229, 0.224, 0.225])
    
    image = Image.open('src/dog.jpg').convert('RGB')
    image_input = image_processor(images=[image], return_tensors="pt")
    with torch.no_grad():
        t_visual_output = torch_model.visual(image_input['pixel_values'])
    
    ## torch_text
    tokenizer = transformers.BertTokenizer.from_pretrained('nghuyong/ernie-3.0-base-zh')

    input_ids = torch.tensor([tokenizer.encode(text="welcome to ernie pytorch project", add_special_tokens=True)])
    with torch.no_grad():
        t_text_output = torch_model.ernie(input_ids)
    print("### pytorch result")
    print("visual_output: ", t_visual_output[:,0][:,:10].numpy())
    print("text_output: ", t_text_output.pooler_output[:,:10].numpy())
    
    ## paddle_visual
    place = "cpu"
    paddle.set_device(place)
    paddle_model = ERNIE_ViL2_base(args)
    paddle_model.eval()
    with paddle.no_grad():
        p_visual_output = paddle_model.visual(paddle.to_tensor(image_input['pixel_values'].numpy()))

    tokenizer = paddlenlp.transformers.AutoTokenizer.from_pretrained("ernie-3.0-base-zh")
    inputs = tokenizer("welcome to ernie pytorch project")
    inputs = {k: paddle.to_tensor([v]) for (k, v) in inputs.items()}
    
    with paddle.no_grad():
        p_text_output = paddle_model.text_model(inputs['input_ids'])
    import pdb; pdb.set_trace()
    print("### paddle result")
    print("visual_output: ", p_visual_output[:,0][:,:10].numpy())
    print("text_output: ", p_text_output[0][:,:10].numpy())
//  [python]
def to_huggingface_processor(self):
        try:
            return ViTFeatureExtractor(
                **{
                    "do_normalize": True,
                    "do_resize": True,
                    "do_rescale": True,
                    "image_mean": list(self.transform.transforms[-1].mean),
                    "image_std": list(self.transform.transforms[-1].std),
                    "resample": PILImageResampling.BILINEAR,
                    "size": self.transform.transforms[0].size,
                }
            )
        except Exception as e:
            print("custom transforms has been provided and is not compatible with this method.")
            print(e)
            return None
//  [python]
def to_huggingface_processor(self):
        try:
            return ViTImageProcessor(
                **{
                    "do_normalize": True,
                    "do_resize": True,
                    "do_rescale": True,
                    "image_mean": list(self.transform.transforms[-1].mean),
                    "image_std": list(self.transform.transforms[-1].std),
                    "resample": PILImageResampling.BILINEAR,
                    "size": self.transform.transforms[0].size,
                }
            )
        except Exception as e:
            print("custom transforms has been provided and is not compatible with this method.")
            print(e)
            return None
//  [python]
def check_ctc_training(self, config, input_values, *args):
        config.ctc_zero_infinity = True
        model = Wav2Vec2ForCTC(config=config)
        model.to(torch_device)
        model.train()

        # freeze feature encoder
        model.freeze_feature_extractor()

        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 lenghts are at least
                # one shorter than logit lenghts 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 check_ctc_training(self, config, input_values, *args):
        config.ctc_zero_infinity = True
        model = Wav2Vec2ForCTC(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 lenghts are at least
                # one shorter than logit lenghts 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 __init__(self, args):
        super().__init__(args)
        self.bert = None
        self.wav2vec2 = None
        self.use_audio_modal = 'audio' in self.args.modality
        self.use_text_modal = 'text' in self.args.modality
        if self.args.audio_auto_config and self.use_audio_modal:
            config = AutoConfig.from_pretrained(self.args.audio_auto_config)
            self.audio_hidden_size = config.hidden_size
            config.mask_feature_prob = args.mask_feature_prob

            self.wav2vec2 = Wav2Vec2ForCTC(config)

            if self.args.freeze_feature_extractor:
                self.wav2vec2.freeze_feature_extractor()

            if self.args.freeze_wav2vec:
                self.freeze_wav2vec_model()

            if self.args.gradient_checkpointing:
                self.wav2vec2.gradient_checkpointing_enable()

        if self.args.text_auto_config and self.use_text_modal:
            config = AutoConfig.from_pretrained(self.args.text_auto_config)
            self.bert_hidden_size = config.hidden_size
            self.bert = AutoModel.from_config(config)

            if self.args.freeze_bert:
                self.freeze_bert_model()

            if self.args.gradient_checkpointing:
                self.bert.gradient_checkpointing_enable()
//  [python]
def __init__(self, args):
        super().__init__(args)
        self.bert = None
        self.wav2vec2 = None
        self.use_audio_modal = 'audio' in self.args.modality
        self.use_text_modal = 'text' in self.args.modality
        if self.args.audio_auto_config and self.use_audio_modal:
            config = AutoConfig.from_pretrained(self.args.audio_auto_config)
            self.audio_hidden_size = config.hidden_size
            config.mask_feature_prob = args.mask_feature_prob

            self.wav2vec2 = Wav2Vec2ForCTC(config)

            if self.args.freeze_feature_extractor:
                self.wav2vec2.freeze_feature_encoder()

            if self.args.freeze_wav2vec:
                self.freeze_wav2vec_model()

            if self.args.gradient_checkpointing:
                self.wav2vec2.gradient_checkpointing_enable()

        if self.args.text_auto_config and self.use_text_modal:
            config = AutoConfig.from_pretrained(self.args.text_auto_config)
            self.bert_hidden_size = config.hidden_size
            self.bert = AutoModel.from_config(config)

            if self.args.freeze_bert:
                self.freeze_bert_model()

            if self.args.gradient_checkpointing:
                self.bert.gradient_checkpointing_enable()
//  [python]
prepared_datasets = speech_datasets.map(prepare_dataset, 
                                 remove_columns=speech_datasets.column_names["train"], 
                                        batch_size=16,
                                        batched=True)

data_collator = DataCollatorCTCWithPadding(processor=processor, padding=True)
wer_metric = load_metric("wer")


# create model
model = Wav2Vec2ForCTC.from_pretrained(
    "facebook/wav2vec2-large-xlsr-53",
    attention_dropout=0.1,
    hidden_dropout=0.1,
    feat_proj_dropout=0.0,
    mask_time_prob=0.05,
    layerdrop=0.1,
    gradient_checkpointing=True,
    ctc_loss_reduction="mean",
    pad_token_id=processor.tokenizer.pad_token_id,
    vocab_size=len(processor.tokenizer)
)
model.freeze_feature_extractor()
//  [python]
prepared_datasets = speech_datasets.map(prepare_dataset, 
                                 remove_columns=speech_datasets.column_names["train"], 
                                        batch_size=16,
                                        batched=True)

data_collator = DataCollatorCTCWithPadding(processor=processor, padding=True)
wer_metric = load_metric("wer")


# create model
model = Wav2Vec2ForCTC.from_pretrained(
    "facebook/wav2vec2-large-xlsr-53",
    attention_dropout=0.1,
    hidden_dropout=0.1,
    feat_proj_dropout=0.0,
    mask_time_prob=0.05,
    layerdrop=0.1,
    gradient_checkpointing=True,
    ctc_loss_reduction="mean",
    pad_token_id=processor.tokenizer.pad_token_id,
    vocab_size=len(processor.tokenizer)
)
model.freeze_feature_encoder()
//  [python]
def check_training(self, config, input_values, *args):
        config.ctc_zero_infinity = True
        model = Wav2Vec2ForCTC(config=config)
        model.to(torch_device)
        model.train()

        # freeze feature encoder
        model.freeze_feature_extractor()

        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 lenghts are at least
                # one shorter than logit lenghts 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 check_training(self, config, input_values, *args):
        config.ctc_zero_infinity = True
        model = Wav2Vec2ForCTC(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 lenghts are at least
                # one shorter than logit lenghts 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]
model =  Wav2Vec2ForCTC.from_pretrained(model_checkpoint) #vocab_size=len(processor.tokenizer)

if hasattr(model, "freeze_feature_extractor"):
  model.freeze_feature_extractor()
//  [python]
model =  Wav2Vec2ForCTC.from_pretrained(model_checkpoint) #vocab_size=len(processor.tokenizer)

if hasattr(model, "freeze_feature_encoder"):
  model.freeze_feature_encoder()
//  [python]
model =  Wav2Vec2ForCTC.from_pretrained(model_checkpoint)
'''model.config.vocab_size = vocab_size
model.config.pad_token_id = processor.tokenizer.pad_token_id
model.lm_head = Linear(in_features=1024, out_features=vocab_size, bias=True)'''

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

if hasattr(model, "freeze_feature_extractor"):
  model.freeze_feature_extractor()
//  [python]
model =  Wav2Vec2ForCTC.from_pretrained(model_checkpoint)
'''model.config.vocab_size = vocab_size
model.config.pad_token_id = processor.tokenizer.pad_token_id
model.lm_head = Linear(in_features=1024, out_features=vocab_size, bias=True)'''

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

if hasattr(model, "freeze_feature_encoder"):
  model.freeze_feature_encoder()
//  [python]
model = Wav2Vec2ForCTC.from_pretrained(
    "facebook/wav2vec2-large-xlsr-53",
    attention_dropout=0.1,
    hidden_dropout=0.1,
    feat_proj_dropout=0.0,
    mask_time_prob=0.05,
    layerdrop=0.1,
    gradient_checkpointing=True,
    ctc_loss_reduction="mean",
    pad_token_id=processor.tokenizer.pad_token_id,
    vocab_size=len(processor.tokenizer)
)

for param in model.parameters():
    param.requires_grad = False
model.freeze_feature_extractor()
//  [python]
model = Wav2Vec2ForCTC.from_pretrained(
    "facebook/wav2vec2-large-xlsr-53",
    attention_dropout=0.1,
    hidden_dropout=0.1,
    feat_proj_dropout=0.0,
    mask_time_prob=0.05,
    layerdrop=0.1,
    gradient_checkpointing=True,
    ctc_loss_reduction="mean",
    pad_token_id=processor.tokenizer.pad_token_id,
    vocab_size=len(processor.tokenizer)
)

for param in model.parameters():
    param.requires_grad = False
model.freeze_feature_encoder()
//  [python]
def train():
    warnings.filterwarnings("ignore", category=np.VisibleDeprecationWarning)
    warnings.filterwarnings("ignore", category=UserWarning)

    with open("config_train.yml") as f:
        args = yaml.load(f, Loader=yaml.FullLoader)

    dataset_train, dataset_test, processor = dataset(args)
    data_collator = DataCollatorCTCWithPadding(processor=processor, padding=True)

    model = Wav2Vec2ForCTC.from_pretrained(
        args["pretrained_checkpoint_dir"],
        attention_dropout=args["attention_dropout"],
        hidden_dropout=args["hidden_dropout"],
        feat_proj_dropout=args["feat_proj_dropout"],
        mask_time_prob=args["mask_time_prob"],
        layerdrop=args["layerdrop"],
        # gradient_checkpointing=args["gradient_checkpointing"],
        ctc_loss_reduction=args["ctc_loss_reduction"],
        pad_token_id=processor.tokenizer.pad_token_id,
        vocab_size=len(processor.tokenizer),
    )
    model.freeze_feature_extractor()

    print("-------load_pretrained_model_done----------")

    training_args = TrainingArguments(
        output_dir=args["checkpoint_dir"],
        group_by_length=args["group_by_length"],
        per_device_train_batch_size=args["batch_size"],
        per_device_eval_batch_size=args["batch_size"],
        gradient_accumulation_steps=args["gradient_accumulation_steps"],
        evaluation_strategy=args["evaluation_strategy"],
        num_train_epochs=args["num_train_epochs"],
        fp16=args["fp16"],
        save_steps=args["save_steps"],
        eval_steps=args["eval_steps"],
        logging_steps=args["logging_steps"],
        weight_decay=args["weight_decay"],
        learning_rate=args["learning_rate"],
        warmup_steps=args["warmup_steps"],
        save_total_limit=args["save_total_limit"],
        dataloader_num_workers=args["dataloader_num_workers"],
    )

    print("-------train_ready_done---------")

    trainer = Trainer(
        model=model,
        data_collator=data_collator,
        args=training_args,
        compute_metrics=compute_metrics,
        train_dataset=dataset_train,
        eval_dataset=dataset_test,
        tokenizer=processor.feature_extractor,
    )

    print("-------training_start!---------")
    trainer.train()
//  [python]
def train():
    warnings.filterwarnings("ignore", category=np.VisibleDeprecationWarning)
    warnings.filterwarnings("ignore", category=UserWarning)

    with open("config_train.yml") as f:
        args = yaml.load(f, Loader=yaml.FullLoader)

    dataset_train, dataset_test, processor = dataset(args)
    data_collator = DataCollatorCTCWithPadding(processor=processor, padding=True)

    model = Wav2Vec2ForCTC.from_pretrained(
        args["pretrained_checkpoint_dir"],
        attention_dropout=args["attention_dropout"],
        hidden_dropout=args["hidden_dropout"],
        feat_proj_dropout=args["feat_proj_dropout"],
        mask_time_prob=args["mask_time_prob"],
        layerdrop=args["layerdrop"],
        # gradient_checkpointing=args["gradient_checkpointing"],
        ctc_loss_reduction=args["ctc_loss_reduction"],
        pad_token_id=processor.tokenizer.pad_token_id,
        vocab_size=len(processor.tokenizer),
    )
    model.freeze_feature_encoder()

    print("-------load_pretrained_model_done----------")

    training_args = TrainingArguments(
        output_dir=args["checkpoint_dir"],
        group_by_length=args["group_by_length"],
        per_device_train_batch_size=args["batch_size"],
        per_device_eval_batch_size=args["batch_size"],
        gradient_accumulation_steps=args["gradient_accumulation_steps"],
        evaluation_strategy=args["evaluation_strategy"],
        num_train_epochs=args["num_train_epochs"],
        fp16=args["fp16"],
        save_steps=args["save_steps"],
        eval_steps=args["eval_steps"],
        logging_steps=args["logging_steps"],
        weight_decay=args["weight_decay"],
        learning_rate=args["learning_rate"],
        warmup_steps=args["warmup_steps"],
        save_total_limit=args["save_total_limit"],
        dataloader_num_workers=args["dataloader_num_workers"],
    )

    print("-------train_ready_done---------")

    trainer = Trainer(
        model=model,
        data_collator=data_collator,
        args=training_args,
        compute_metrics=compute_metrics,
        train_dataset=dataset_train,
        eval_dataset=dataset_test,
        tokenizer=processor.feature_extractor,
    )

    print("-------training_start!---------")
    trainer.train()
//  [python]
def test_processor_case_2(self):
        # case 2: document image classification (training, inference) + token classification (inference), apply_ocr=False

        feature_extractor = LayoutLMv2FeatureExtractor(apply_ocr=False)
        tokenizers = self.get_tokenizers
        images = self.get_images

        for tokenizer in tokenizers:
            processor = LayoutXLMProcessor(feature_extractor=feature_extractor, tokenizer=tokenizer)

            # not batched
            words = ["hello", "world"]
            boxes = [[1, 2, 3, 4], [5, 6, 7, 8]]
            input_processor = processor(images[0], words, boxes=boxes, return_tensors="pt")

            # verify keys
            expected_keys = ["input_ids", "bbox", "attention_mask", "image"]
            actual_keys = list(input_processor.keys())
            for key in expected_keys:
                self.assertIn(key, actual_keys)

            # verify input_ids
            expected_decoding = "
 hello world
"
            decoding = tokenizer.decode(input_processor.input_ids.squeeze().tolist())
            self.assertSequenceEqual(decoding, expected_decoding)

            # batched
            words = [["hello", "world"], ["my", "name", "is", "niels"]]
            boxes = [[[1, 2, 3, 4], [5, 6, 7, 8]], [[3, 2, 5, 1], [6, 7, 4, 2], [3, 9, 2, 4], [1, 1, 2, 3]]]
            input_processor = processor(images, words, boxes=boxes, padding=True, return_tensors="pt")

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids"]
            actual_keys = sorted(list(input_processor.keys()))
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
            expected_decoding = "
 hello world
<pad><pad>"
            decoding = tokenizer.decode(input_processor.input_ids[0].tolist())
            self.assertSequenceEqual(decoding, expected_decoding)

            # verify bbox
            expected_bbox = [
                [0, 0, 0, 0],
                [3, 2, 5, 1],
                [6, 7, 4, 2],
                [3, 9, 2, 4],
                [1, 1, 2, 3],
                [1, 1, 2, 3],
                [1000, 1000, 1000, 1000],
            ]
            self.assertListEqual(input_processor.bbox[1].tolist(), expected_bbox)
//  [python]
def test_processor_case_2(self):
        # case 2: document image classification (training, inference) + token classification (inference), apply_ocr=False

        feature_extractor = LayoutLMv2ImageProcessor(apply_ocr=False)
        tokenizers = self.get_tokenizers
        images = self.get_images

        for tokenizer in tokenizers:
            processor = LayoutXLMProcessor(feature_extractor=feature_extractor, tokenizer=tokenizer)

            # not batched
            words = ["hello", "world"]
            boxes = [[1, 2, 3, 4], [5, 6, 7, 8]]
            input_processor = processor(images[0], words, boxes=boxes, return_tensors="pt")

            # verify keys
            expected_keys = ["input_ids", "bbox", "attention_mask", "image"]
            actual_keys = list(input_processor.keys())
            for key in expected_keys:
                self.assertIn(key, actual_keys)

            # verify input_ids
            expected_decoding = "
 hello world
"
            decoding = tokenizer.decode(input_processor.input_ids.squeeze().tolist())
            self.assertSequenceEqual(decoding, expected_decoding)

            # batched
            words = [["hello", "world"], ["my", "name", "is", "niels"]]
            boxes = [[[1, 2, 3, 4], [5, 6, 7, 8]], [[3, 2, 5, 1], [6, 7, 4, 2], [3, 9, 2, 4], [1, 1, 2, 3]]]
            input_processor = processor(images, words, boxes=boxes, padding=True, return_tensors="pt")

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids"]
            actual_keys = sorted(list(input_processor.keys()))
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
            expected_decoding = "
 hello world
<pad><pad>"
            decoding = tokenizer.decode(input_processor.input_ids[0].tolist())
            self.assertSequenceEqual(decoding, expected_decoding)

            # verify bbox
            expected_bbox = [
                [0, 0, 0, 0],
                [3, 2, 5, 1],
                [6, 7, 4, 2],
                [3, 9, 2, 4],
                [1, 1, 2, 3],
                [1, 1, 2, 3],
                [1000, 1000, 1000, 1000],
            ]
            self.assertListEqual(input_processor.bbox[1].tolist(), expected_bbox)
//  [python]
def test_processor_case_3(self):
        # case 3: token classification (training), apply_ocr=False

        feature_extractor = LayoutLMv2FeatureExtractor(apply_ocr=False)
        tokenizers = self.get_tokenizers
        images = self.get_images

        for tokenizer in tokenizers:
            processor = LayoutXLMProcessor(feature_extractor=feature_extractor, tokenizer=tokenizer)

            # not batched
            words = ["weirdly", "world"]
            boxes = [[1, 2, 3, 4], [5, 6, 7, 8]]
            word_labels = [1, 2]
            input_processor = processor(images[0], words, boxes=boxes, word_labels=word_labels, return_tensors="pt")

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids", "labels"]
            actual_keys = sorted(list(input_processor.keys()))
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
            expected_decoding = "
 weirdly world
"
            decoding = tokenizer.decode(input_processor.input_ids.squeeze().tolist())
            self.assertSequenceEqual(decoding, expected_decoding)

            # verify labels
            expected_labels = [-100, 1, -100, 2, -100]
            self.assertListEqual(input_processor.labels.squeeze().tolist(), expected_labels)

            # batched
            words = [["hello", "world"], ["my", "name", "is", "niels"]]
            boxes = [[[1, 2, 3, 4], [5, 6, 7, 8]], [[3, 2, 5, 1], [6, 7, 4, 2], [3, 9, 2, 4], [1, 1, 2, 3]]]
            word_labels = [[1, 2], [6, 3, 10, 2]]
            input_processor = processor(
                images, words, boxes=boxes, word_labels=word_labels, padding=True, return_tensors="pt"
            )

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids", "labels"]
            actual_keys = sorted(list(input_processor.keys()))
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
            expected_decoding = "
 my name is niels
"
            decoding = tokenizer.decode(input_processor.input_ids[1].tolist())
            self.assertSequenceEqual(decoding, expected_decoding)

            # verify bbox
            expected_bbox = [
                [0, 0, 0, 0],
                [3, 2, 5, 1],
                [6, 7, 4, 2],
                [3, 9, 2, 4],
                [1, 1, 2, 3],
                [1, 1, 2, 3],
                [1000, 1000, 1000, 1000],
            ]
            self.assertListEqual(input_processor.bbox[1].tolist(), expected_bbox)

            # verify labels
            expected_labels = [-100, 6, 3, 10, 2, -100, -100]
            self.assertListEqual(input_processor.labels[1].tolist(), expected_labels)
//  [python]
def test_processor_case_3(self):
        # case 3: token classification (training), apply_ocr=False

        feature_extractor = LayoutLMv2ImageProcessor(apply_ocr=False)
        tokenizers = self.get_tokenizers
        images = self.get_images

        for tokenizer in tokenizers:
            processor = LayoutXLMProcessor(feature_extractor=feature_extractor, tokenizer=tokenizer)

            # not batched
            words = ["weirdly", "world"]
            boxes = [[1, 2, 3, 4], [5, 6, 7, 8]]
            word_labels = [1, 2]
            input_processor = processor(images[0], words, boxes=boxes, word_labels=word_labels, return_tensors="pt")

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids", "labels"]
            actual_keys = sorted(list(input_processor.keys()))
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
            expected_decoding = "
 weirdly world
"
            decoding = tokenizer.decode(input_processor.input_ids.squeeze().tolist())
            self.assertSequenceEqual(decoding, expected_decoding)

            # verify labels
            expected_labels = [-100, 1, -100, 2, -100]
            self.assertListEqual(input_processor.labels.squeeze().tolist(), expected_labels)

            # batched
            words = [["hello", "world"], ["my", "name", "is", "niels"]]
            boxes = [[[1, 2, 3, 4], [5, 6, 7, 8]], [[3, 2, 5, 1], [6, 7, 4, 2], [3, 9, 2, 4], [1, 1, 2, 3]]]
            word_labels = [[1, 2], [6, 3, 10, 2]]
            input_processor = processor(
                images, words, boxes=boxes, word_labels=word_labels, padding=True, return_tensors="pt"
            )

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids", "labels"]
            actual_keys = sorted(list(input_processor.keys()))
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
            expected_decoding = "
 my name is niels
"
            decoding = tokenizer.decode(input_processor.input_ids[1].tolist())
            self.assertSequenceEqual(decoding, expected_decoding)

            # verify bbox
            expected_bbox = [
                [0, 0, 0, 0],
                [3, 2, 5, 1],
                [6, 7, 4, 2],
                [3, 9, 2, 4],
                [1, 1, 2, 3],
                [1, 1, 2, 3],
                [1000, 1000, 1000, 1000],
            ]
            self.assertListEqual(input_processor.bbox[1].tolist(), expected_bbox)

            # verify labels
            expected_labels = [-100, 6, 3, 10, 2, -100, -100]
            self.assertListEqual(input_processor.labels[1].tolist(), expected_labels)
//  [python]
def test_processor_case_5(self):
        # case 5: visual question answering (inference), apply_ocr=False

        feature_extractor = LayoutLMv2FeatureExtractor(apply_ocr=False)
        tokenizers = self.get_tokenizers
        images = self.get_images

        for tokenizer in tokenizers:
            processor = LayoutXLMProcessor(feature_extractor=feature_extractor, tokenizer=tokenizer)

            # not batched
            question = "What's his name?"
            words = ["hello", "world"]
            boxes = [[1, 2, 3, 4], [5, 6, 7, 8]]
            input_processor = processor(images[0], question, words, boxes, return_tensors="pt")

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids"]
            actual_keys = sorted(list(input_processor.keys()))
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
            expected_decoding = "
 What's his name?

 hello world
"
            decoding = tokenizer.decode(input_processor.input_ids.squeeze().tolist())
            self.assertSequenceEqual(decoding, expected_decoding)

            # batched
            questions = ["How old is he?", "what's the time"]
            words = [["hello", "world"], ["my", "name", "is", "niels"]]
            boxes = [[[1, 2, 3, 4], [5, 6, 7, 8]], [[3, 2, 5, 1], [6, 7, 4, 2], [3, 9, 2, 4], [1, 1, 2, 3]]]
            input_processor = processor(images, questions, words, boxes, padding=True, return_tensors="pt")

            # verify keys
            expected_keys = ["attention_mask", "bbox", "image", "input_ids"]
            actual_keys = sorted(list(input_processor.keys()))
            self.assertListEqual(actual_keys, expected_keys)

            # verify input_ids
            expected_decoding = "
 How old is he?

 hello world
<pad><pad>"
            decoding = tokenizer.decode(input_processor.input_ids[0].tolist())
            self.assertSequenceEqual(decoding, expected_decoding)

            expected_decoding = "
 what's the time

 my name is niels
"
            decoding = tokenizer.decode(input_processor.input_ids[1].tolist())
            self.assertSequenceEqual(decoding, expected_decoding)

            # verify bbox
            expected_bbox = [[6, 7, 4, 2], [3, 9, 2, 4], [1, 1, 2, 3], [1, 1, 2, 3], [1000, 1000, 1000, 1000]]
            self.assertListEqual(input_processor.bbox[1].tolist()[-5:], expected_bbox)