Tamás Ficsor

add model

11cdb73 over 1 year ago

No virus

14.9 kB

	import math
	from math import gcd
	import functools
	import torch
	import torch.nn.functional as F
	from torch import nn, einsum

	from einops import rearrange, reduce, repeat
	from einops.layers.torch import Rearrange
	from transformers.modeling_utils import PreTrainedModel


	def exists(val):
	return val is not None


	def lcm(*numbers):
	return int(functools.reduce(lambda x, y: int((x * y) / gcd(x, y)), numbers, 1))


	def masked_mean(tensor, mask, dim = -1):
	diff_len = len(tensor.shape) - len(mask.shape)
	mask = mask[(..., ((None,) diff_len))]
	tensor.masked_fill_(~mask, 0.)

	total_el = mask.sum(dim = dim)
	mean = tensor.sum(dim = dim) / total_el.clamp(min = 1.)
	mean.masked_fill_(total_el == 0, 0.)
	return mean


	def next_divisible_length(seqlen, multiple):
	return math.ceil(seqlen / multiple) * multiple


	def pad_to_multiple(tensor, multiple, *, seq_dim, dim = -1, value = 0.):
	seqlen = tensor.shape[seq_dim]
	length = next_divisible_length(seqlen, multiple)
	if length == seqlen:
	return tensor
	remainder = length - seqlen
	pad_offset = (0,) * (-1 - dim) * 2
	return F.pad(tensor, (*pad_offset, 0, remainder), value = value)


	# helper classes
	class Pad(nn.Module):
	def __init__(self, padding, value = 0.):
	super().__init__()
	self.padding = padding
	self.value = value

	def forward(self, x):
	return F.pad(x, self.padding, value = self.value)


	class DepthwiseConv1d(nn.Module):
	def __init__(self, dim_in, dim_out, kernel_size):
	super().__init__()
	self.conv = nn.Conv1d(dim_in, dim_out, kernel_size, groups = dim_in)
	self.proj_out = nn.Conv1d(dim_out, dim_out, 1)

	def forward(self, x):
	x = self.conv(x)
	return self.proj_out(x)


	# main class
	class GBST(PreTrainedModel):
	def _init_weights(self, module):
	"""Initialize the weights"""
	if isinstance(module, nn.Linear):
	# Slightly different from the TF version which uses truncated_normal for initialization
	# cf https://github.com/pytorch/pytorch/pull/5617
	module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
	if module.bias is not None:
	module.bias.data.zero_()
	elif isinstance(module, nn.Embedding):
	module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
	if module.padding_idx is not None:
	module.weight.data[module.padding_idx].zero_()
	elif isinstance(module, nn.LayerNorm):
	module.bias.data.zero_()
	module.weight.data.fill_(1.0)

	def __init__(
	self,
	*,
	num_tokens,
	dim,
	max_block_size = None,
	blocks = None,
	downsample_factor = 4,
	score_consensus_attn = True,
	return_without_downsample = True,
	config = None
	):
	super(GBST, self).__init__(config=config)
	assert exists(max_block_size) ^ exists(blocks), 'either max_block_size or blocks are given on initialization'
	self.word_embeddings = nn.Embedding(num_tokens, dim)
	self.position_embeddings = nn.Embedding(config.max_position_embeddings, dim)
	self.token_type_embeddings = nn.Embedding(config.type_vocab_size, dim)

	self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))

	self.return_without_downsample = return_without_downsample

	if exists(blocks):
	assert isinstance(blocks, tuple), 'blocks must be a tuple of block sizes'
	self.blocks = tuple(map(lambda el: el if isinstance(el, tuple) else (el, 0), blocks))
	assert all([(offset < block_size) for block_size, offset in self.blocks]), 'offset must be always smaller than the block size'

	max_block_size = max(list(map(lambda t: t[0], self.blocks)))
	else:
	self.blocks = tuple(map(lambda el: (el, 0), range(1, max_block_size + 1)))

	self.pos_conv = nn.Sequential(
	Pad((0, 0, 0, max_block_size - 1)),
	Rearrange('b n d -> b d n'),
	DepthwiseConv1d(dim, dim, kernel_size = max_block_size),
	Rearrange('b d n -> b n d')
	)

	self.score_fn = nn.Sequential(
	nn.Linear(dim, 1),
	Rearrange('... () -> ...')
	)

	self.score_consensus_attn = score_consensus_attn

	assert downsample_factor <= max_block_size, 'final downsample factor should be less than the maximum block size'

	self.block_pad_multiple = lcm(*[block_size for block_size, _ in self.blocks])
	self.downsample_factor = downsample_factor

	def forward(self, input_ids, attention_mask=None, position_ids=None, token_type_ids=None, inputs_embeds=None):
	b, n, block_mult, ds_factor, device = *input_ids.shape, self.block_pad_multiple, self.downsample_factor, input_ids.device
	m = next_divisible_length(n, ds_factor)

	# get character token embeddings

	input_ids = self.word_embeddings(input_ids)
	token_type_embeddings = self.token_type_embeddings(token_type_ids)

	seq_len = input_ids.size()[1]
	position_ids = self.position_ids[:, :seq_len]
	position_embeddings = self.position_embeddings(position_ids)

	input_ids = input_ids + token_type_embeddings + position_embeddings
	# do a conv to generate the positions for the tokens

	input_ids = self.pos_conv(input_ids)

	# pad both sequence and attention_mask to length visibile by all block sizes from 0 to max block size

	input_ids = pad_to_multiple(input_ids, block_mult, seq_dim=1, dim=-2)

	if exists(attention_mask):
	attention_mask = pad_to_multiple(attention_mask, block_mult, seq_dim=1, dim=-1, value=False)

	# compute representations for all blocks by mean pooling

	block_masks = []
	block_reprs = []

	for block_size, offset in self.blocks:
	# clone the input sequence as well as the attention_mask, in order to pad for offsets

	block_x = input_ids.clone()

	if exists(attention_mask):
	block_mask = attention_mask.clone()

	# pad for offsets, if needed

	need_padding = offset > 0

	if need_padding:
	left_offset, right_offset = (block_size - offset), offset
	block_x = F.pad(block_x, (0, 0, left_offset, right_offset), value = 0.)

	if exists(attention_mask):
	block_mask = F.pad(block_mask, (left_offset, right_offset), value = False)

	# group input sequence into blocks

	blocks = rearrange(block_x, 'b (n m) d -> b n m d', m = block_size)

	# either mean pool the blocks, or do a masked mean

	if exists(attention_mask):
	mask_blocks = rearrange(block_mask, 'b (n m) -> b n m', m = block_size)
	block_repr = masked_mean(blocks, mask_blocks, dim = -2)
	else:
	block_repr = blocks.mean(dim = -2)

	# append the block representations, as well as the pooled block masks

	block_repr = repeat(block_repr, 'b n d -> b (n m) d', m = block_size)

	if need_padding:
	block_repr = block_repr[:, left_offset:-right_offset]

	block_reprs.append(block_repr)

	if exists(attention_mask):
	mask_blocks = torch.any(mask_blocks, dim = -1)
	mask_blocks = repeat(mask_blocks, 'b n -> b (n m)', m = block_size)

	if need_padding:
	mask_blocks = mask_blocks[:, left_offset:-right_offset]

	block_masks.append(mask_blocks)

	# stack all the block representations

	block_reprs = torch.stack(block_reprs, dim = 2)

	# calculate scores and softmax across the block size dimension

	scores = self.score_fn(block_reprs)

	if exists(attention_mask):
	block_masks = torch.stack(block_masks, dim = 2)
	max_neg_value = -torch.finfo(scores.dtype).max
	scores = scores.masked_fill(~block_masks, max_neg_value)

	scores = scores.softmax(dim = 2)

	# do the cheap consensus attention, eq (5) in paper

	if self.score_consensus_attn:
	score_sim = einsum('b i d, b j d -> b i j', scores, scores)

	if exists(attention_mask):
	cross_mask = rearrange(attention_mask, 'b i -> b i ()') * rearrange(attention_mask, 'b j -> b () j')
	max_neg_value = -torch.finfo(score_sim.dtype).max
	score_sim = score_sim.masked_fill(~cross_mask, max_neg_value)

	score_attn = score_sim.softmax(dim=-1)
	scores = einsum('b i j, b j m -> b i m', score_attn, scores)

	# multiply the block representations by the position-wise scores

	scores = rearrange(scores, 'b n m -> b n m ()')
	input_ids = (block_reprs * scores).sum(dim=2)

	# truncate to length divisible by downsample factor

	input_ids = input_ids[:, :m]

	original = None
	if self.return_without_downsample:
	original = torch.clone(input_ids)

	input_ids, attention_mask = self.down_sample(input_ids, attention_mask, ds_factor)

	return input_ids, attention_mask, original

	@staticmethod
	def down_sample(input_ids, attention_mask, ds_factor):
	n = input_ids.shape[1]
	m = next_divisible_length(n, ds_factor)
	if exists(attention_mask):
	attention_mask = attention_mask[:, :m]

	# final mean pooling downsample
	input_ids = rearrange(input_ids, 'b (n m) d -> b n m d', m=ds_factor)

	if exists(attention_mask):
	attention_mask = rearrange(attention_mask, 'b (n m) -> b n m', m=ds_factor)
	input_ids = masked_mean(input_ids, attention_mask, dim=2)
	attention_mask = torch.any(attention_mask, dim=-1)
	else:
	input_ids = input_ids.mean(dim=-2)
	return input_ids, attention_mask

	def block_score(self, input_ids, attention_mask=None, position_ids=None, token_type_ids=None, inputs_embeds=None):
	b, n, block_mult, ds_factor, device = *input_ids.shape, self.block_pad_multiple, self.downsample_factor, input_ids.device
	m = next_divisible_length(n, ds_factor)

	# get character token embeddings

	input_ids = self.word_embeddings(input_ids)

	# do a conv to generate the positions for the tokens

	input_ids = self.pos_conv(input_ids)

	# pad both sequence and attention_mask to length visibile by all block sizes from 0 to max block size

	input_ids = pad_to_multiple(input_ids, block_mult, seq_dim=1, dim=-2)

	if exists(attention_mask):
	attention_mask = pad_to_multiple(attention_mask, block_mult, seq_dim=1, dim=-1, value=False)

	# compute representations for all blocks by mean pooling

	block_masks = []
	block_reprs = []

	for block_size, offset in self.blocks:
	# clone the input sequence as well as the attention_mask, in order to pad for offsets

	block_x = input_ids.clone()

	if exists(attention_mask):
	block_mask = attention_mask.clone()

	# pad for offsets, if needed

	need_padding = offset > 0

	if need_padding:
	left_offset, right_offset = (block_size - offset), offset
	block_x = F.pad(block_x, (0, 0, left_offset, right_offset), value = 0.)

	if exists(attention_mask):
	block_mask = F.pad(block_mask, (left_offset, right_offset), value = False)

	# group input sequence into blocks

	blocks = rearrange(block_x, 'b (n m) d -> b n m d', m = block_size)

	# either mean pool the blocks, or do a masked mean

	if exists(attention_mask):
	mask_blocks = rearrange(block_mask, 'b (n m) -> b n m', m = block_size)
	block_repr = masked_mean(blocks, mask_blocks, dim = -2)
	else:
	block_repr = blocks.mean(dim = -2)

	# append the block representations, as well as the pooled block masks

	block_repr = repeat(block_repr, 'b n d -> b (n m) d', m = block_size)

	if need_padding:
	block_repr = block_repr[:, left_offset:-right_offset]

	block_reprs.append(block_repr)

	if exists(attention_mask):
	mask_blocks = torch.any(mask_blocks, dim = -1)
	mask_blocks = repeat(mask_blocks, 'b n -> b (n m)', m = block_size)

	if need_padding:
	mask_blocks = mask_blocks[:, left_offset:-right_offset]

	block_masks.append(mask_blocks)

	# stack all the block representations

	block_reprs = torch.stack(block_reprs, dim = 2)

	# calculate scores and softmax across the block size dimension

	scores = self.score_fn(block_reprs)

	if exists(attention_mask):
	block_masks = torch.stack(block_masks, dim = 2)
	max_neg_value = -torch.finfo(scores.dtype).max
	scores = scores.masked_fill(~block_masks, max_neg_value)

	scores = scores.softmax(dim = 2)

	# do the cheap consensus attention, eq (5) in paper

	if self.score_consensus_attn:
	score_sim = einsum('b i d, b j d -> b i j', scores, scores)

	if exists(attention_mask):
	cross_mask = rearrange(attention_mask, 'b i -> b i ()') * rearrange(attention_mask, 'b j -> b () j')
	max_neg_value = -torch.finfo(score_sim.dtype).max
	score_sim = score_sim.masked_fill(~cross_mask, max_neg_value)

	score_attn = score_sim.softmax(dim=-1)
	scores = einsum('b i j, b j m -> b i m', score_attn, scores)

	# multiply the block representations by the position-wise scores

	scores = rearrange(scores, 'b n m -> b n m ()')
	input_ids = (block_reprs * scores).sum(dim=2)

	# truncate to length divisible by downsample factor

	input_ids = input_ids[:, :m]

	if exists(attention_mask):
	attention_mask = attention_mask[:, :m]

	original = None
	if self.return_without_downsample:
	original = torch.clone(input_ids)

	# final mean pooling downsample
	input_ids = rearrange(input_ids, 'b (n m) d -> b n m d', m=ds_factor)

	if exists(attention_mask):
	attention_mask = rearrange(attention_mask, 'b (n m) -> b n m', m=ds_factor)
	input_ids = masked_mean(input_ids, attention_mask, dim=2)
	attention_mask = torch.any(attention_mask, dim=-1)
	else:
	input_ids = input_ids.mean(dim=-2)

	return scores