Spaces:
Running
Running
File size: 2,369 Bytes
e76b79a | 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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | # Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from typing import List, Optional
import torch
import torch.nn.functional as F
def get_line_info():
from inspect import getframeinfo, currentframe
caller = getframeinfo(currentframe().f_back)
return f"{caller.filename}:{caller.lineno}"
def assert_shape(tensor, *shape_args):
# This only returns the unbound shapes
source = f'{get_line_info()}:'
assert len(tensor.shape) == len(shape_args), f"{source} should have {len(shape_args)} dimensions, actually has shape {tensor.shape}"
unbound = []
for i, ii in enumerate(shape_args):
if ii == -1:
unbound.append(tensor.shape[i])
else:
assert (
tensor.shape[i] == ii
), f"{source} shape should be {shape_args}, actually {tensor.shape}"
if len(unbound) == 1:
return unbound[0]
return unbound
def assert_probs(tensor: torch.FloatTensor):
""" Assert that [*,K] input tensor is valid probs in final dim. """
assert is_prob_tensor(tensor)
def assert_logprobs(tensor: torch.FloatTensor):
""" Assert that [*,K] input tensor is valid logprobs in final dim. """
try:
assert_probs(tensor.exp())
except AssertionError:
raise AssertionError('Not logprobs, perhaps you have logits and need to apply F.log_softmax?')
def is_1hot_tensor(x1h):
return x1h.min() == 0 and x1h.max() == 1 and \
x1h.shape[-1] > 1 and (x1h.sum(-1) == 1).all()
def is_prob_tensor(p, atol=1e-3):
return (
p.min() >= 0
and p.max() <= 1
and torch.isclose(
p.sum(axis=-1), torch.ones_like(p.sum(axis=-1)), atol=atol).all()
).item()
def add_eos_bos(
seq1h,
bos_idx: Optional[int] = None,
eos_idx: Optional[int] = None,
):
"""
Helper for (possibly) prepending bos/cls and appending eos tokens.
"""
B, L, K = seq1h.shape
to_concat = []
if bos_idx is not None:
to_concat.append(F.one_hot(torch.full([B,1], bos_idx), K).to(seq1h))
to_concat.append(seq1h)
if eos_idx is not None:
to_concat.append(F.one_hot(torch.full([B,1], eos_idx), K).to(seq1h))
seq1h = torch.cat(to_concat, axis=1)
return seq1h
|