wehub-resource-sync
888 lines
32 KiB
Python
888 lines
32 KiB
Python
# coding=utf-8
|
|
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
|
|
# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# Copyright (c) Facebook, Inc. All Rights Reserved
|
|
|
|
|
|
import torch
|
|
import math
|
|
import pickle
|
|
import random
|
|
import os
|
|
import numpy as np
|
|
|
|
from collections import deque
|
|
from typing import Optional, Tuple, List
|
|
from .processor import (
|
|
Processor,
|
|
MetaProcessor,
|
|
TextProcessor,
|
|
Aligner,
|
|
MMAttentionMask2DProcessor
|
|
)
|
|
|
|
from ..utils import ShardedTensor
|
|
|
|
|
|
class How2MetaProcessor(MetaProcessor):
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
path = self._get_split_path(config)
|
|
with open(path) as fd:
|
|
self.data = [line.strip() for line in fd]
|
|
|
|
def __getitem__(self, idx):
|
|
video_id = self.data[idx]
|
|
return video_id, video_id
|
|
|
|
|
|
class ShardedHow2MetaProcessor(How2MetaProcessor):
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
self.split = str(config.split)
|
|
self.vfeat_dir = config.vfeat_dir
|
|
self._init_shard()
|
|
|
|
def _init_shard(self):
|
|
if self.split == "train":
|
|
meta_fn = os.path.join(self.vfeat_dir, "train" + "_meta.pkl")
|
|
with open(meta_fn, "rb") as fr:
|
|
meta = pickle.load(fr)
|
|
elif self.split == "valid":
|
|
meta_fn = os.path.join(self.vfeat_dir, "val" + "_meta.pkl")
|
|
with open(meta_fn, "rb") as fr:
|
|
meta = pickle.load(fr)
|
|
elif self.split == "test":
|
|
print("use how2 val as test.")
|
|
meta_fn = os.path.join(self.vfeat_dir, "val" + "_meta.pkl")
|
|
with open(meta_fn, "rb") as fr:
|
|
meta = pickle.load(fr)
|
|
else:
|
|
raise ValueError("unsupported for MetaProcessor:", self.split)
|
|
video_id_to_shard = {}
|
|
for shard_id in meta:
|
|
for video_idx, video_id in enumerate(meta[shard_id]):
|
|
video_id_to_shard[video_id] = (shard_id, video_idx)
|
|
self.video_id_to_shard = video_id_to_shard
|
|
|
|
def __getitem__(self, idx):
|
|
video_id, video_id = super().__getitem__(idx)
|
|
shard_id, shard_idx = self.video_id_to_shard[video_id]
|
|
meta = (video_id, idx, shard_id, shard_idx)
|
|
return meta, meta
|
|
|
|
|
|
class ShardedVideoProcessor(Processor):
|
|
"""
|
|
mmaped shards of numpy video features.
|
|
"""
|
|
|
|
def __init__(self, config):
|
|
self.split = str(config.split)
|
|
self.vfeat_dir = config.vfeat_dir
|
|
|
|
def __call__(self, video_id):
|
|
_, _, shard_id, video_idx = video_id
|
|
if self.split == "train":
|
|
shard = ShardedTensor.load(
|
|
os.path.join(self.vfeat_dir, "train" + "_" + str(shard_id)),
|
|
"r"
|
|
)
|
|
elif self.split == "valid":
|
|
shard = ShardedTensor.load(
|
|
os.path.join(self.vfeat_dir, "val" + "_" + str(shard_id)),
|
|
"r"
|
|
)
|
|
elif self.split == "test":
|
|
shard = ShardedTensor.load(
|
|
os.path.join(self.vfeat_dir, "val" + "_" + str(shard_id)),
|
|
"r"
|
|
)
|
|
else:
|
|
raise ValueError("unknown split", self.split)
|
|
feat = shard[video_idx]
|
|
return feat
|
|
|
|
|
|
class ShardedTextProcessor(Processor):
|
|
def __init__(self, config):
|
|
self.tfeat_dir = str(config.tfeat_dir)
|
|
self.split = str(config.split)
|
|
|
|
def __call__(self, video_id):
|
|
_, _, shard_id, shard_idx = video_id
|
|
if self.split == "train":
|
|
target_path = self.tfeat_dir + "train" + "_" + str(shard_id)
|
|
elif self.split == "valid":
|
|
target_path = self.tfeat_dir + "val" + "_" + str(shard_id)
|
|
elif self.split == "test":
|
|
target_path = self.tfeat_dir + "val" + "_" + str(shard_id)
|
|
else:
|
|
raise ValueError("unknown split", self.split)
|
|
|
|
startend = ShardedTensor.load(
|
|
target_path + ".startends", "r")[shard_idx]
|
|
cap_ids = ShardedTensor.load(
|
|
target_path + ".caps_ids", "r")[shard_idx]
|
|
cap = []
|
|
for clip_idx in range(len(cap_ids)):
|
|
clip = cap_ids[clip_idx]
|
|
cap.append(clip[clip != -1].tolist())
|
|
start, end = startend[:, 0].tolist(), startend[:, 1].tolist()
|
|
return {"start": start, "end": end, "cap": cap}
|
|
|
|
|
|
class FixedLenAligner(Aligner):
|
|
"""
|
|
In the model we assume text is on the left (closer to BERT formulation)
|
|
and video is on the right.
|
|
We fix the total length of text + video.
|
|
max_video_len is in number of secs.
|
|
max_text_len is in number of tokens.
|
|
|
|
special tokens formats:
|
|
we use the format [CLS] [SEP] text tokens [SEP] [PAD] ...
|
|
[CLS] will be splitted out into:
|
|
[CLS] video tokens [SEP] text tokens [SEP] [PAD] ...
|
|
token_type_ids will be generated by the model (for now).
|
|
0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
|
|
| first sequence | second sequence |
|
|
so each sequence owns a [SEP] token for no-ops.
|
|
"""
|
|
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
self.text_clip_sampler = TextClipSamplingProcessor(
|
|
self.max_len - self.max_video_len - 3
|
|
)
|
|
"""
|
|
decide subsampling:
|
|
`config.subsampling` will change batch_size in trainer.
|
|
`config.clip_per_video` (used by RetriTask) doesn't
|
|
change batch_size in trainer.
|
|
"""
|
|
subsampling = config.subsampling \
|
|
if config.subsampling is not None else None
|
|
if config.clip_per_video is not None:
|
|
subsampling = config.clip_per_video
|
|
self.subsampling = subsampling
|
|
|
|
def _get_text_maxlen(self):
|
|
# use max text len
|
|
return self.text_clip_sampler.max_text_len
|
|
|
|
def __call__(self, video_id, video_feature, text_feature):
|
|
from transformers import default_data_collator
|
|
video_idx = video_id[1]
|
|
if self.subsampling is not None and self.subsampling >= 1:
|
|
batch = []
|
|
for _ in range(self.subsampling):
|
|
centerclip_idx = random.randint(
|
|
0, len(text_feature["start"]) - 1)
|
|
batch.append(
|
|
self.sampling(
|
|
video_idx,
|
|
video_feature,
|
|
text_feature,
|
|
centerclip_idx,
|
|
self._get_text_maxlen()
|
|
))
|
|
batch = self.batch_post_processing(batch, video_feature)
|
|
batch = default_data_collator(batch)
|
|
else:
|
|
raise ValueError(
|
|
"dataset.subsampling must be >= 1 for efficient video loading.")
|
|
batch = self.sampling(video_idx, video_feature, text_feature)
|
|
batch = self.batch_post_processing(batch, video_feature)
|
|
|
|
batch["video_id"] = video_id if isinstance(video_id, str) \
|
|
else video_id[0]
|
|
# e2e: make sure frame ids is into tensor.
|
|
assert torch.is_tensor(batch["vfeats"])
|
|
return batch
|
|
|
|
def sampling(
|
|
self,
|
|
video_idx,
|
|
video_feature,
|
|
text_feature,
|
|
centerclip_idx=None,
|
|
sampled_max_text_len=None,
|
|
):
|
|
text_clip_indexs = self.text_clip_sampler(
|
|
text_feature, centerclip_idx,
|
|
sampled_max_text_len
|
|
)
|
|
if isinstance(video_feature, np.ndarray):
|
|
video_len = len(video_feature)
|
|
else:
|
|
video_len = math.ceil(text_feature["end"][-1])
|
|
|
|
video_end = min(
|
|
math.ceil(text_feature["end"][text_clip_indexs[-1]]),
|
|
video_len
|
|
)
|
|
video_start = max(
|
|
min(
|
|
math.floor(text_feature["start"][text_clip_indexs[0]]),
|
|
video_end),
|
|
0
|
|
)
|
|
|
|
video_clips = {"start": [video_start], "end": [video_end]}
|
|
|
|
# tensorize.
|
|
vfeats, vmasks = self._build_video_seq(
|
|
video_feature, video_clips
|
|
)
|
|
caps, cmasks = self._build_text_seq(
|
|
text_feature, text_clip_indexs
|
|
)
|
|
|
|
text_start = text_clip_indexs[0]
|
|
text_end = text_clip_indexs[-1] + 1
|
|
|
|
return {
|
|
"caps": caps,
|
|
"cmasks": cmasks,
|
|
"vfeats": vfeats,
|
|
"vmasks": vmasks,
|
|
"video_start": video_start,
|
|
"video_end": video_end,
|
|
"text_start": text_start,
|
|
"text_end": text_end,
|
|
}
|
|
|
|
|
|
class VariedLenAligner(FixedLenAligner):
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
self.sampled_min_len = config.sampled_min_len
|
|
self.sampled_max_len = config.sampled_max_len
|
|
|
|
def _get_text_maxlen(self):
|
|
return random.randint(self.sampled_min_len, self.sampled_max_len)
|
|
|
|
|
|
class StartClipAligner(VariedLenAligner):
|
|
def sampling(
|
|
self,
|
|
video_idx,
|
|
video_feature,
|
|
text_feature,
|
|
centerclip_idx=None,
|
|
sampled_max_text_len=None,
|
|
):
|
|
return super().sampling(
|
|
video_idx, video_feature, text_feature, 0)
|
|
|
|
|
|
class OverlappedAligner(VariedLenAligner):
|
|
"""video clip and text clip has overlappings
|
|
but may not be the same start/end."""
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
self.sampled_video_min_len = config.sampled_video_min_len
|
|
self.sampled_video_max_len = config.sampled_video_max_len
|
|
|
|
self.video_clip_sampler = VideoClipSamplingProcessor()
|
|
|
|
def _get_video_maxlen(self):
|
|
return random.randint(
|
|
self.sampled_video_min_len, self.sampled_video_max_len)
|
|
|
|
def sampling(
|
|
self,
|
|
video_idx,
|
|
video_feature,
|
|
text_feature,
|
|
centerclip_idx=None,
|
|
sampled_max_text_len=None,
|
|
):
|
|
text_clip_indexs = self.text_clip_sampler(
|
|
text_feature, centerclip_idx,
|
|
sampled_max_text_len
|
|
)
|
|
if isinstance(video_feature, np.ndarray):
|
|
video_len = len(video_feature)
|
|
else:
|
|
video_len = math.ceil(text_feature["end"][-1])
|
|
low = math.floor(text_feature["start"][text_clip_indexs[0]])
|
|
high = math.ceil(text_feature["end"][text_clip_indexs[-1]])
|
|
if low < high:
|
|
center = random.randint(low, high)
|
|
else:
|
|
center = int((low + high) // 2)
|
|
center = max(0, min(video_feature.shape[0] - 1, center))
|
|
|
|
assert 0 <= center < video_feature.shape[0]
|
|
|
|
video_clips = self.video_clip_sampler(
|
|
video_len, self._get_video_maxlen(), center
|
|
)
|
|
video_start = video_clips["start"][0]
|
|
video_end = video_clips["end"][0]
|
|
|
|
# tensorize.
|
|
vfeats, vmasks = self._build_video_seq(
|
|
video_feature, video_clips
|
|
)
|
|
caps, cmasks = self._build_text_seq(
|
|
text_feature, text_clip_indexs
|
|
)
|
|
|
|
text_start = text_clip_indexs[0]
|
|
text_end = text_clip_indexs[-1] + 1
|
|
|
|
return {
|
|
"caps": caps,
|
|
"cmasks": cmasks,
|
|
"vfeats": vfeats,
|
|
"vmasks": vmasks,
|
|
"video_start": video_start,
|
|
"video_end": video_end,
|
|
"text_start": text_start,
|
|
"text_end": text_end,
|
|
}
|
|
|
|
|
|
class MFMMLMAligner(FixedLenAligner):
|
|
"""
|
|
`FixedLenAligner` with Masked Language Model and Masked Frame Model.
|
|
"""
|
|
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
keep_prob = config.keep_prob if config.keep_prob is not None else 1.0
|
|
self.text_clip_sampler = TextClipSamplingProcessor(
|
|
self.max_len - self.max_video_len - 3, keep_prob
|
|
)
|
|
self.sampled_min_len = config.sampled_min_len
|
|
self.sampled_max_len = config.sampled_max_len
|
|
self.masked_token_sampler = TextMaskingProcessor(config)
|
|
self.mm_type = config.mm_type \
|
|
if config.mm_type is not None else "full"
|
|
self.attnmasker = MMAttentionMask2DProcessor() \
|
|
if self.mm_type == "textgen" else None
|
|
self.masked_frame_sampler = FrameMaskingProcessor(config)
|
|
self.lazy_vfeat_mask = (
|
|
False if config.lazy_vfeat_mask is None else config.lazy_vfeat_mask
|
|
)
|
|
self.mm_prob = config.mm_prob if config.mm_prob is not None else 0.
|
|
|
|
def __call__(self, video_id, video_feature, text_feature):
|
|
from transformers import default_data_collator
|
|
if self.subsampling is not None and self.subsampling > 1:
|
|
batch = []
|
|
for _ in range(self.subsampling):
|
|
centerclip_idx = random.randint(
|
|
0, len(text_feature["start"]) - 1)
|
|
sampled_max_text_len = random.randint(
|
|
self.sampled_min_len, self.sampled_max_len
|
|
)
|
|
batch.append(
|
|
self.sampling(
|
|
video_id,
|
|
video_feature,
|
|
text_feature,
|
|
centerclip_idx,
|
|
sampled_max_text_len,
|
|
)
|
|
)
|
|
batch = self.batch_post_processing(batch, video_feature)
|
|
batch = default_data_collator(batch)
|
|
else:
|
|
batch = self.sampling(video_id, video_feature, text_feature)
|
|
batch = self.batch_post_processing(batch, video_feature)
|
|
batch["video_id"] = video_id if isinstance(video_id, str) \
|
|
else video_id[0]
|
|
return batch
|
|
|
|
def sampling(
|
|
self,
|
|
video_id,
|
|
video_feature,
|
|
text_feature,
|
|
centerclip_idx=None,
|
|
sampled_max_text_len=None,
|
|
):
|
|
output = FixedLenAligner.sampling(self,
|
|
video_id, video_feature, text_feature,
|
|
centerclip_idx, sampled_max_text_len)
|
|
|
|
masking_text, masking_video = None, None
|
|
if random.random() < self.mm_prob:
|
|
if random.random() > 0.5:
|
|
masking_text, masking_video = self.mm_type, "no"
|
|
else:
|
|
masking_text, masking_video = "no", "full"
|
|
video_feats = output["vfeats"] if not self.lazy_vfeat_mask else None
|
|
video_label = self.masked_frame_sampler(
|
|
output["vmasks"], masking_video, vfeats=video_feats)
|
|
caps, text_label = self.masked_token_sampler(
|
|
output["caps"], masking_text)
|
|
|
|
output.update({
|
|
"caps": caps,
|
|
"video_label": video_label,
|
|
"text_label": text_label,
|
|
})
|
|
|
|
if self.attnmasker is not None:
|
|
attention_mask = self.attnmasker(
|
|
output["vmasks"], output["cmasks"], masking_text)
|
|
output.update({
|
|
"attention_mask": attention_mask
|
|
})
|
|
return output
|
|
|
|
|
|
class FrameMaskingProcessor(Processor):
|
|
def __init__(self, config):
|
|
self.mfm_probability = 0.15
|
|
if config.mfm_probability is not None:
|
|
self.mfm_probability = config.mfm_probability
|
|
|
|
def __call__(self, vmasks, modality_masking=None, vfeats=None):
|
|
"""
|
|
We perform lazy masking to save data transfer time.
|
|
It only generates video_labels by default and MFM model
|
|
will do actualy masking.
|
|
Return: `video_label` is a binary mask.
|
|
"""
|
|
video_label = vmasks.clone()
|
|
if modality_masking is not None:
|
|
if modality_masking == "full":
|
|
probability_matrix = torch.full(video_label.shape, 1.)
|
|
elif modality_masking == "no":
|
|
probability_matrix = torch.full(video_label.shape, 0.)
|
|
elif modality_masking == "inverse":
|
|
probability_matrix = torch.full(
|
|
video_label.shape, 1. - self.mfm_probability)
|
|
else:
|
|
raise ValueError("unknown modality masking.", modality_masking)
|
|
else:
|
|
probability_matrix = torch.full(
|
|
video_label.shape, self.mfm_probability)
|
|
masked_indices = torch.bernoulli(probability_matrix).bool()
|
|
# We only compute loss on masked tokens
|
|
video_label[~masked_indices] = 0
|
|
if vfeats is not None:
|
|
vfeats[video_label, :] = 0.0
|
|
return video_label
|
|
|
|
|
|
class TextGenerationProcessor(Processor):
|
|
def __init__(self, tokenizer):
|
|
self.bos_token_id = tokenizer.bos_token_id
|
|
self.pad_token_id = tokenizer.pad_token_id
|
|
|
|
def __call__(self, inputs):
|
|
labels = inputs.clone()
|
|
# [CLS] [SEP] for video
|
|
labels[:2] = -100
|
|
# keep [SEP] for text.
|
|
pad_mask = labels == self.pad_token_id
|
|
labels[pad_mask] = -100
|
|
inputs[2:] = torch.cat([
|
|
torch.LongTensor([self.bos_token_id]),
|
|
inputs[2:-1]])
|
|
inputs[pad_mask] = self.pad_token_id
|
|
assert len(inputs) == len(labels)
|
|
return inputs, labels
|
|
|
|
|
|
class TextMaskingProcessor(Processor):
|
|
def __init__(self, config):
|
|
"""this function is borrowed from
|
|
`transformers/data/data_collator.DataCollatorForLanguageModeling`"""
|
|
self.mlm_probability = 0.15
|
|
if config.mlm_probability is not None:
|
|
self.mlm_probability = config.mlm_probability
|
|
self.bert_name = config.bert_name
|
|
# [CLS] is used as bos_token and [SEP] is used as eos_token.
|
|
# https://huggingface.co/transformers/master/model_doc/bertgeneration.html
|
|
from transformers import AutoTokenizer
|
|
self.tokenizer = AutoTokenizer.from_pretrained(
|
|
self.bert_name, bos_token="[CLS]", eos_token="[SEP]")
|
|
self.textgen = TextGenerationProcessor(self.tokenizer)
|
|
|
|
def __call__(
|
|
self, inputs: torch.Tensor,
|
|
modality_masking=None,
|
|
special_tokens_mask: Optional[torch.Tensor] = None
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
expand modality_masking into
|
|
None: traditional bert masking.
|
|
"no": no masking.
|
|
"full": all [MASK] token for generation.
|
|
"gen": autoregressive generation.
|
|
"""
|
|
"""
|
|
Prepare masked tokens inputs/labels for masked language modeling:
|
|
80% MASK, 10% random, 10% original.
|
|
"""
|
|
labels = inputs.clone()
|
|
# We sample a few tokens in each sequence for MLM training
|
|
# (with probability `self.mlm_probability`)
|
|
if modality_masking is not None:
|
|
if modality_masking == "full":
|
|
probability_matrix = torch.full(labels.shape, 1.)
|
|
elif modality_masking == "no":
|
|
probability_matrix = torch.full(labels.shape, 0.)
|
|
elif modality_masking.startswith("textgen"):
|
|
# [CLS] [SEP] <s> ...
|
|
inputs, labels = self.textgen(inputs)
|
|
if "mask" not in modality_masking:
|
|
return inputs, labels
|
|
inputs = self.mask_input(inputs, special_tokens_mask)
|
|
return inputs, labels
|
|
elif modality_masking == "mask":
|
|
inputs = self.mask_input(inputs, special_tokens_mask)
|
|
labels = torch.full(inputs.shape, -100)
|
|
return inputs, labels
|
|
elif modality_masking == "inverse":
|
|
probability_matrix = torch.full(labels.shape, 1. - self.mlm_probability)
|
|
else:
|
|
raise ValueError("unknown modality masking.", modality_masking)
|
|
else:
|
|
probability_matrix = torch.full(labels.shape, self.mlm_probability)
|
|
|
|
if special_tokens_mask is None:
|
|
special_tokens_mask = self.get_special_tokens_mask(
|
|
labels.tolist(), already_has_special_tokens=True
|
|
)
|
|
special_tokens_mask = torch.tensor(
|
|
special_tokens_mask, dtype=torch.bool)
|
|
else:
|
|
special_tokens_mask = special_tokens_mask.bool()
|
|
|
|
probability_matrix.masked_fill_(special_tokens_mask, value=0.0)
|
|
masked_indices = torch.bernoulli(probability_matrix).bool()
|
|
labels[~masked_indices] = -100 # We only compute loss on masked tokens
|
|
|
|
# 80% of the time,
|
|
# we replace masked input tokens with tokenizer.mask_token ([MASK])
|
|
indices_replaced = (
|
|
torch.bernoulli(
|
|
torch.full(labels.shape, 0.8)).bool() & masked_indices
|
|
)
|
|
inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(
|
|
self.tokenizer.mask_token
|
|
)
|
|
|
|
# 10% of the time, we replace masked input tokens with random word
|
|
indices_random = (
|
|
torch.bernoulli(torch.full(labels.shape, 0.5)).bool()
|
|
& masked_indices
|
|
& ~indices_replaced
|
|
)
|
|
random_words = torch.randint(
|
|
len(self.tokenizer), labels.shape, dtype=torch.long
|
|
)
|
|
inputs[indices_random] = random_words[indices_random]
|
|
|
|
# The rest of the time (10% of the time) we keep the masked input
|
|
# tokens unchanged
|
|
return inputs, labels
|
|
|
|
def mask_input(self, inputs, special_tokens_mask=None):
|
|
# the following is new with masked autoregressive.
|
|
probability_matrix = torch.full(
|
|
inputs.shape, self.mlm_probability)
|
|
if special_tokens_mask is None:
|
|
special_tokens_mask = self.get_special_tokens_mask(
|
|
inputs.tolist(), already_has_special_tokens=True
|
|
)
|
|
special_tokens_mask = torch.tensor(
|
|
special_tokens_mask, dtype=torch.bool)
|
|
else:
|
|
special_tokens_mask = special_tokens_mask.bool()
|
|
probability_matrix.masked_fill_(special_tokens_mask, value=0.0)
|
|
masked_indices = torch.bernoulli(probability_matrix).bool()
|
|
indices_replaced = (
|
|
torch.bernoulli(
|
|
torch.full(inputs.shape, 0.8)).bool() & masked_indices
|
|
)
|
|
inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(
|
|
self.tokenizer.mask_token
|
|
)
|
|
|
|
# 10% of the time, we replace masked input tokens with random word
|
|
indices_random = (
|
|
torch.bernoulli(torch.full(inputs.shape, 0.5)).bool()
|
|
& masked_indices
|
|
& ~indices_replaced
|
|
)
|
|
random_words = torch.randint(
|
|
len(self.tokenizer), inputs.shape, dtype=torch.long
|
|
)
|
|
inputs[indices_random] = random_words[indices_random]
|
|
return inputs
|
|
|
|
def get_special_tokens_mask(
|
|
self, token_ids_0: List[int],
|
|
token_ids_1: Optional[List[int]] = None,
|
|
already_has_special_tokens: bool = False
|
|
) -> List[int]:
|
|
"""
|
|
Note: the version from transformers do not consider pad
|
|
as special tokens.
|
|
"""
|
|
|
|
if already_has_special_tokens:
|
|
if token_ids_1 is not None:
|
|
raise ValueError(
|
|
"You should not supply a second sequence if"
|
|
"the provided sequence of "
|
|
"ids is already formated with special tokens "
|
|
"for the model."
|
|
)
|
|
return list(map(lambda x: 1 if x in [
|
|
self.tokenizer.sep_token_id,
|
|
self.tokenizer.cls_token_id,
|
|
self.tokenizer.pad_token_id] else 0, token_ids_0))
|
|
|
|
if token_ids_1 is not None:
|
|
return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
|
|
return [1] + ([0] * len(token_ids_0)) + [1]
|
|
|
|
|
|
class TextClipSamplingProcessor(Processor):
|
|
def __init__(self, max_text_len, keep_prob=1.0):
|
|
self.max_text_len = max_text_len
|
|
self.max_video_len = 256 # always hold.
|
|
self.keep_prob = keep_prob
|
|
|
|
def __call__(
|
|
self,
|
|
text_feature,
|
|
centerclip_idx=None,
|
|
sampled_max_text_len=None,
|
|
sampled_max_video_len=None,
|
|
):
|
|
# Let's use all caps for now and see if 256 can cover all of them.
|
|
if sampled_max_text_len is not None:
|
|
max_text_len = sampled_max_text_len
|
|
else:
|
|
max_text_len = self.max_text_len
|
|
if sampled_max_video_len is not None:
|
|
max_video_len = sampled_max_video_len
|
|
else:
|
|
max_video_len = self.max_video_len
|
|
|
|
t_num_clips = len(text_feature["start"])
|
|
|
|
if centerclip_idx is None:
|
|
centerclip_idx = random.randint(0, t_num_clips - 1)
|
|
|
|
start_idx, end_idx = centerclip_idx, centerclip_idx + 1
|
|
text_clip_indexs = deque()
|
|
text_clip_indexs.append(start_idx)
|
|
text_len = len(text_feature["cap"][start_idx])
|
|
|
|
video_len = max(
|
|
0,
|
|
text_feature["end"][start_idx]
|
|
- text_feature["start"][start_idx],
|
|
)
|
|
|
|
while (
|
|
(start_idx > 0 or end_idx < t_num_clips)
|
|
and text_len < max_text_len
|
|
and video_len < max_video_len
|
|
):
|
|
if random.random() > 0.5 and end_idx < t_num_clips:
|
|
# skip the next one?
|
|
if random.random() > self.keep_prob and (end_idx + 1) < t_num_clips:
|
|
end_idx = end_idx + 1
|
|
text_clip_indexs.append(end_idx)
|
|
text_len += len(text_feature["cap"][end_idx])
|
|
end_idx += 1
|
|
elif start_idx > 0:
|
|
if random.random() > self.keep_prob and (start_idx - 1) > 0:
|
|
start_idx = start_idx - 1
|
|
start_idx -= 1
|
|
text_clip_indexs.insert(0, start_idx)
|
|
text_len += len(text_feature["cap"][start_idx])
|
|
else:
|
|
if end_idx < t_num_clips:
|
|
if random.random() > self.keep_prob and (end_idx + 1) < t_num_clips:
|
|
end_idx = end_idx + 1
|
|
text_clip_indexs.append(end_idx)
|
|
text_len += len(text_feature["cap"][end_idx])
|
|
end_idx += 1
|
|
else:
|
|
return text_clip_indexs
|
|
video_len = max(
|
|
0,
|
|
text_feature["end"][text_clip_indexs[-1]]
|
|
- text_feature["start"][text_clip_indexs[0]],
|
|
)
|
|
return text_clip_indexs
|
|
|
|
|
|
class VideoClipSamplingProcessor(Processor):
|
|
def __call__(self, video_len, max_video_len, center):
|
|
"""
|
|
`video_len`: length of the video.
|
|
`max_video_len`: maximum video tokens allowd in a sequence.
|
|
`center`: initial starting index.
|
|
"""
|
|
assert center >= 0 and center < video_len
|
|
t_clip_len = 0
|
|
start, end = center, center
|
|
while (start > 0 or end < video_len) and t_clip_len < max_video_len:
|
|
# decide the direction to grow.
|
|
if start <= 0:
|
|
end += 1
|
|
elif end >= video_len:
|
|
start -= 1
|
|
elif random.random() > 0.5:
|
|
end += 1
|
|
else:
|
|
start -= 1
|
|
t_clip_len += 1
|
|
return {"start": [start], "end": [end]}
|
|
|
|
|
|
class How2MILNCEAligner(FixedLenAligner):
|
|
"""reference: `antoine77340/MIL-NCE_HowTo100M/video_loader.py`"""
|
|
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
self.num_candidates = 4
|
|
self.min_time = 5.0
|
|
self.num_sec = 3.2
|
|
# self.num_sec = self.num_frames / float(self.fps) num_frames=16 / fps = 5
|
|
# self.num_frames = 16
|
|
|
|
def sampling(
|
|
self,
|
|
video_id,
|
|
video_feature,
|
|
text_feature,
|
|
centerclip_idx=None, # will be ignored.
|
|
sampled_max_text_len=None # will be ignored.
|
|
):
|
|
text, start, end = self._get_text(text_feature)
|
|
video = self._get_video(video_feature, start, end)
|
|
|
|
vfeats = torch.zeros((self.max_video_len, video_feature.shape[1]))
|
|
vmasks = torch.zeros((self.max_video_len,), dtype=torch.bool)
|
|
vfeats[: video.shape[0]] = torch.from_numpy(np.array(video))
|
|
vmasks[: video.shape[0]] = 1
|
|
|
|
caps, cmasks = [], []
|
|
for words in text:
|
|
cap, cmask = self._build_text_seq(text_feature, words)
|
|
caps.append(cap)
|
|
cmasks.append(cmask)
|
|
caps = torch.stack(caps)
|
|
cmasks = torch.stack(cmasks)
|
|
# video of shape: (video_len)
|
|
# text of shape (num_candidates, max_text_len)
|
|
|
|
return {
|
|
"caps": caps,
|
|
"cmasks": cmasks,
|
|
"vfeats": vfeats,
|
|
"vmasks": vmasks,
|
|
# "video_id": video_id,
|
|
}
|
|
|
|
def _get_video(self, video_feature, start, end):
|
|
start_seek = random.randint(start, int(max(start, end - self.num_sec)))
|
|
# duration = self.num_sec + 0.1
|
|
return video_feature[start_seek : int(start_seek + self.num_sec)]
|
|
|
|
def _get_text(self, cap):
|
|
ind = random.randint(0, len(cap["start"]) - 1)
|
|
if self.num_candidates == 1:
|
|
words = [ind]
|
|
else:
|
|
words = []
|
|
cap_start = self._find_nearest_candidates(cap, ind)
|
|
for i in range(self.num_candidates):
|
|
words.append([max(0, min(len(cap["cap"]) - 1, cap_start + i))])
|
|
|
|
start, end = cap["start"][ind], cap["end"][ind]
|
|
# TODO: May need to be improved for edge cases.
|
|
# expand the min time.
|
|
if end - start < self.min_time:
|
|
diff = self.min_time - end + start
|
|
start = max(0, start - diff / 2)
|
|
end = start + self.min_time
|
|
return words, int(start), int(end)
|
|
|
|
def _find_nearest_candidates(self, caption, ind):
|
|
"""find the range of the clips."""
|
|
start, end = ind, ind
|
|
#diff = caption["end"][end] - caption["start"][start]
|
|
n_candidate = 1
|
|
while n_candidate < self.num_candidates:
|
|
# the first clip
|
|
if start == 0:
|
|
return 0
|
|
# we add () in the following condition to fix the bug.
|
|
elif end == (len(caption["start"]) - 1):
|
|
return start - (self.num_candidates - n_candidate)
|
|
elif (caption["end"][end] - caption["start"][start - 1]) < (
|
|
caption["end"][end + 1] - caption["start"][start]
|
|
):
|
|
start -= 1
|
|
else:
|
|
end += 1
|
|
n_candidate += 1
|
|
return start
|
|
|
|
|
|
class PKLJSONStrTextProcessor(TextProcessor):
|
|
"""`caption.json` from howto100m are preprocessed as a
|
|
dict `[video_id, json_str]`.
|
|
Json parsing tokenization are conducted on-the-fly and cached into dict.
|
|
"""
|
|
|
|
def __init__(self, config, max_clip_text_len=96):
|
|
print("[Warning] PKLJSONStrTextProcessor is slow for num_workers > 0.")
|
|
self.caption_pkl_path = str(config.caption_pkl_path)
|
|
with open(self.caption_pkl_path, "rb") as fd:
|
|
self.data = pickle.load(fd)
|
|
self.max_clip_text_len = max_clip_text_len
|
|
from transformers import AutoTokenizer
|
|
self.tokenizer = AutoTokenizer.from_pretrained(
|
|
str(config.bert_name), use_fast=config.use_fast
|
|
)
|
|
|
|
def __call__(self, video_id):
|
|
caption = self.data[video_id]
|
|
if isinstance(caption, str):
|
|
import json
|
|
caption = json.loads(caption)
|
|
cap = []
|
|
for clip_idx, text_clip in enumerate(caption["text"]):
|
|
clip_ids = []
|
|
if isinstance(text_clip, str):
|
|
clip_ids = self.tokenizer(
|
|
text_clip[: self.max_clip_text_len],
|
|
add_special_tokens=False
|
|
)["input_ids"]
|
|
cap.append(clip_ids)
|
|
caption["cap"] = cap
|
|
caption.pop("text") # save space.
|
|
self.data[video_id] = caption
|
|
return caption
|