# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/00-core.ipynb (unless otherwise specified).

__all__ = ['flatten_dict', 'stack_dicts', 'add_suffix', 'pad_to_size', 'logprobs_from_logits', 'whiten',

'clip_by_value', 'entropy_from_logits', 'average_torch_dicts', 'stats_to_np', 'build_bert_batch_from_txt']

# Cell

import torch

import torch.nn.functional as F

import collections

import numpy as np

# Cell

def flatten_dict(nested, sep='/'):

"""Flatten dictionary and concatenate nested keys with separator."""

def rec(nest, prefix, into):

for k, v in nest.items():

if sep in k:

raise ValueError(f"separator '{sep}' not allowed to be in key '{k}'")

if isinstance(v, collections.Mapping):

rec(v, prefix + k + sep, into)

else:

into[prefix + k] = v

flat = {}

rec(nested, '', flat)

return flat

def stack_dicts(stats_dicts):

"""Stack the values of a dict."""

results = dict()

for k in stats_dicts[0]:

stats_list = [torch.flatten(d[k]) for d in stats_dicts]

results[k] = torch.stack(stats_list)

return results

def add_suffix(input_dict, suffix):

"""Add suffix to dict keys."""

return dict((k + suffix, v) for k,v in input_dict.items())

# Cell

def pad_to_size(tensor, size, dim=1, padding=50256):

"""Pad tensor to size."""

t_size = tensor.size()[dim]

if t_size==size:

return tensor

else:

return torch.nn.functional.pad(tensor, (0,size-t_size), 'constant', padding)

def logprobs_from_logits(logits, labels):

"""

See: https://github.com/pytorch/pytorch/issues/563#issuecomment-330103591

"""

logp = F.log_softmax(logits, dim=2)

logpy = torch.gather(logp, 2, labels.unsqueeze(2)).squeeze(-1)

return logpy

def whiten(values, shift_mean=True):

"""Whiten values."""

mean, var = torch.mean(values), torch.var(values)

whitened = (values - mean) * torch.rsqrt(var + 1e-8)

if not shift_mean:

whitened += mean

return whitened

def clip_by_value(x, tensor_min, tensor_max):

"""

Tensor extenstion to torch.clamp

https://github.com/pytorch/pytorch/issues/2793#issuecomment-428784713

"""

clipped = torch.max(torch.min(x, tensor_max), tensor_min)

return clipped

def entropy_from_logits(logits):

"""Calculate entropy from logits."""

pd = torch.nn.functional.softmax(logits, dim=-1)

entropy = torch.logsumexp(logits, axis=-1) - torch.sum(pd*logits, axis=-1)

return entropy

def average_torch_dicts(list_of_dicts):

"""Average values of a list of dicts wiht torch tensors."""

average_dict = dict()

for key in list_of_dicts[0].keys():

average_dict[key] = torch.mean(torch.stack([d[key] for d in list_of_dicts]), axis=0)

return average_dict

def stats_to_np(stats_dict):

"""Cast all torch.tensors in dict to numpy arrays."""

new_dict = dict()

for k, v in stats_dict.items():

if isinstance(v, torch.Tensor):

new_dict[k] = v.detach().cpu().numpy()

else:

new_dict[k] = v

if np.isscalar(new_dict[k]):

new_dict[k] = float(new_dict[k])

return new_dict

# Cell

def build_bert_batch_from_txt(text_list, tokenizer, device):

"""Create token id and attention mask tensors from text list for BERT classification."""

# tokenize

tensors = [tokenizer.encode(txt, return_tensors="pt").to(device) for txt in text_list]

# find max length to pad to

max_len = max([t.size()[1] for t in tensors])

# get padded tensors and attention masks

# (attention masks make bert ignore padding)

padded_tensors = []

attention_masks = []

for tensor in tensors:

attention_mask = torch.ones(tensor.size(), device=device)

padded_tensors.append(pad_to_size(tensor, max_len, padding=0))

attention_masks.append(pad_to_size(attention_mask, max_len, padding=0))

# stack all tensors

padded_tensors = torch.cat(padded_tensors)

attention_masks = torch.cat(attention_masks)

return padded_tensors, attention_masks