File size: 1,209 Bytes
988f575 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | import torch
def postprocess_text(preds, labels):
preds = [pred.strip() for pred in preds]
labels = [[label.strip()] for label in labels]
return preds, labels
# Add collate_fn to DataLoader
def collate_fn(batch):
# Add padding to the inputs
# "inputs" must be 250 frames long
# "attention_mask" must be 250 frames long
# "labels" must be 128 tokens long
return {
"sign_inputs": torch.stack([
torch.cat((sample["sign_inputs"], torch.zeros(250 - sample["sign_inputs"].shape[0], 208)), dim=0)
for sample in batch
]),
"attention_mask": torch.stack([
torch.cat((sample["attention_mask"], torch.zeros(250 - sample["attention_mask"].shape[0])), dim=0)
if sample["attention_mask"].shape[0] < 250
else sample["attention_mask"]
for sample in batch
]),
"labels": torch.stack([
torch.cat((sample["labels"].squeeze(0), torch.zeros(128 - sample["labels"].shape[0])), dim=0)
if sample["labels"].shape[0] < 128
else sample["labels"]
for sample in batch
]).squeeze(0).to(torch.long),
}
|