|
|
import torch
|
|
|
from safetensors.torch import load_file
|
|
|
|
|
|
from torch.nn.utils import remove_weight_norm
|
|
|
import warnings
|
|
|
warnings.simplefilter(action='ignore', category=FutureWarning)
|
|
|
|
|
|
|
|
|
def load_ckpt_state_dict(ckpt_path):
|
|
|
if ckpt_path.endswith(".safetensors"):
|
|
|
state_dict = load_file(ckpt_path)
|
|
|
else:
|
|
|
state_dict = torch.load(ckpt_path, map_location="cpu")["state_dict"]
|
|
|
|
|
|
return state_dict
|
|
|
|
|
|
def remove_weight_norm_from_model(model):
|
|
|
for module in model.modules():
|
|
|
if hasattr(module, "weight"):
|
|
|
print(f"Removing weight norm from {module}")
|
|
|
remove_weight_norm(module)
|
|
|
|
|
|
return model
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def multinomial(input: torch.Tensor, num_samples: int, replacement=False, *, generator=None):
|
|
|
"""torch.multinomial with arbitrary number of dimensions, and number of candidates on the last dimension.
|
|
|
|
|
|
Args:
|
|
|
input (torch.Tensor): The input tensor containing probabilities.
|
|
|
num_samples (int): Number of samples to draw.
|
|
|
replacement (bool): Whether to draw with replacement or not.
|
|
|
Keywords args:
|
|
|
generator (torch.Generator): A pseudorandom number generator for sampling.
|
|
|
Returns:
|
|
|
torch.Tensor: Last dimension contains num_samples indices
|
|
|
sampled from the multinomial probability distribution
|
|
|
located in the last dimension of tensor input.
|
|
|
"""
|
|
|
|
|
|
if num_samples == 1:
|
|
|
q = torch.empty_like(input).exponential_(1, generator=generator)
|
|
|
return torch.argmax(input / q, dim=-1, keepdim=True).to(torch.int64)
|
|
|
|
|
|
input_ = input.reshape(-1, input.shape[-1])
|
|
|
output_ = torch.multinomial(input_, num_samples=num_samples, replacement=replacement, generator=generator)
|
|
|
output = output_.reshape(*list(input.shape[:-1]), -1)
|
|
|
return output
|
|
|
|
|
|
|
|
|
def sample_top_k(probs: torch.Tensor, k: int) -> torch.Tensor:
|
|
|
"""Sample next token from top K values along the last dimension of the input probs tensor.
|
|
|
|
|
|
Args:
|
|
|
probs (torch.Tensor): Input probabilities with token candidates on the last dimension.
|
|
|
k (int): The k in “top-k”.
|
|
|
Returns:
|
|
|
torch.Tensor: Sampled tokens.
|
|
|
"""
|
|
|
top_k_value, _ = torch.topk(probs, k, dim=-1)
|
|
|
min_value_top_k = top_k_value[..., [-1]]
|
|
|
probs *= (probs >= min_value_top_k).float()
|
|
|
probs.div_(probs.sum(dim=-1, keepdim=True))
|
|
|
next_token = multinomial(probs, num_samples=1)
|
|
|
return next_token
|
|
|
|
|
|
|
|
|
def sample_top_p(probs: torch.Tensor, p: float) -> torch.Tensor:
|
|
|
"""Sample next token from top P probabilities along the last dimension of the input probs tensor.
|
|
|
|
|
|
Args:
|
|
|
probs (torch.Tensor): Input probabilities with token candidates on the last dimension.
|
|
|
p (int): The p in “top-p”.
|
|
|
Returns:
|
|
|
torch.Tensor: Sampled tokens.
|
|
|
"""
|
|
|
probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
|
|
|
probs_sum = torch.cumsum(probs_sort, dim=-1)
|
|
|
mask = probs_sum - probs_sort > p
|
|
|
probs_sort *= (~mask).float()
|
|
|
probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
|
|
|
next_token = multinomial(probs_sort, num_samples=1)
|
|
|
next_token = torch.gather(probs_idx, -1, next_token)
|
|
|
return next_token
|
|
|
|
|
|
def next_power_of_two(n):
|
|
|
return 2 ** (n - 1).bit_length()
|
|
|
|
|
|
def next_multiple_of_64(n):
|
|
|
return ((n + 63) // 64) * 64 |
