Spaces:
Build error
Build error
| """rotary_positional_embedding.py - Rotary Positional Embedding | |
| code from github.com/lucidrains/rotary-embedding-torch | |
| MIT License | |
| """ | |
| from math import pi, log | |
| import torch | |
| from torch import nn, einsum | |
| from einops import rearrange, repeat | |
| def exists(val): | |
| return val is not None | |
| def broadcat(tensors, dim=-1): | |
| num_tensors = len(tensors) | |
| shape_lens = set(list(map(lambda t: len(t.shape), tensors))) | |
| assert len(shape_lens) == 1, 'tensors must all have the same number of dimensions' | |
| shape_len = list(shape_lens)[0] | |
| dim = (dim + shape_len) if dim < 0 else dim | |
| dims = list(zip(*map(lambda t: list(t.shape), tensors))) | |
| expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim] | |
| assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims) | |
| ]), 'invalid dimensions for broadcastable concatentation' | |
| max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims)) | |
| expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims)) | |
| expanded_dims.insert(dim, (dim, dims[dim])) | |
| expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims))) | |
| tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes))) | |
| return torch.cat(tensors, dim=dim) | |
| # rotary embedding helper functions | |
| def rotate_half(x): | |
| x = rearrange(x, '... (d r) -> ... d r', r=2) | |
| x1, x2 = x.unbind(dim=-1) | |
| x = torch.stack((-x2, x1), dim=-1) | |
| return rearrange(x, '... d r -> ... (d r)') | |
| def apply_rotary_emb(freqs, t, start_index=0, scale=1.): | |
| rot_dim, seq_len = freqs.shape[-1], t.shape[-2] | |
| freqs = freqs[-seq_len:, :] | |
| freqs = freqs.to(t) | |
| end_index = start_index + rot_dim | |
| assert rot_dim <= t.shape[ | |
| -1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}' | |
| t_left, t, t_right = t[..., :start_index], t[..., start_index:end_index], t[..., end_index:] | |
| t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale) | |
| return torch.cat((t_left, t, t_right), dim=-1) | |
| # learned rotation helpers | |
| def apply_learned_rotations(rotations, t, start_index=0, freq_ranges=None): | |
| if exists(freq_ranges): | |
| rotations = einsum('..., f -> ... f', rotations, freq_ranges) | |
| rotations = rearrange(rotations, '... r f -> ... (r f)') | |
| rotations = repeat(rotations, '... n -> ... (n r)', r=2) | |
| return apply_rotary_emb(rotations, t, start_index=start_index) | |
| # classes | |
| class RotaryEmbedding(nn.Module): | |
| def __init__(self, | |
| dim, | |
| custom_freqs=None, | |
| freqs_for='lang', | |
| theta=10000, | |
| max_freq=10, | |
| num_freqs=1, | |
| learned_freq=False, | |
| use_xpos=False, | |
| xpos_scale_base=512, | |
| interpolate_factor=1., | |
| theta_rescale_factor=1.): | |
| super().__init__() | |
| # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning | |
| # has some connection to NTK literature | |
| # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/ | |
| theta *= theta_rescale_factor**(dim / (dim - 2)) | |
| if exists(custom_freqs): | |
| freqs = custom_freqs | |
| elif freqs_for == 'lang': | |
| freqs = 1. / (theta**(torch.arange(0, dim, 2)[:(dim // 2)].float() / dim)) | |
| elif freqs_for == 'pixel': | |
| freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi | |
| elif freqs_for == 'constant': | |
| freqs = torch.ones(num_freqs).float() | |
| else: | |
| raise ValueError(f'unknown modality {freqs_for}') | |
| self.cache = dict() | |
| self.cache_scale = dict() | |
| self.freqs = nn.Parameter(freqs, requires_grad=learned_freq) | |
| # interpolation factors | |
| assert interpolate_factor >= 1. | |
| self.interpolate_factor = interpolate_factor | |
| # xpos | |
| self.use_xpos = use_xpos | |
| if not use_xpos: | |
| self.register_buffer('scale', None) | |
| return | |
| scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim) | |
| self.scale_base = xpos_scale_base | |
| self.register_buffer('scale', scale) | |
| def get_seq_pos(self, seq_len, device, dtype, offset=0): | |
| return (torch.arange(seq_len, device=device, dtype=dtype) + | |
| offset) / self.interpolate_factor | |
| def rotate_queries_or_keys(self, t, seq_dim=-2, offset=0, freq_seq_len=None): | |
| assert not self.use_xpos, 'you must use `.rotate_queries_and_keys` method instead and pass in both queries and keys, for length extrapolatable rotary embeddings' | |
| device, dtype, seq_len = t.device, t.dtype, t.shape[seq_dim] | |
| if exists(freq_seq_len): | |
| assert freq_seq_len >= seq_len | |
| seq_len = freq_seq_len | |
| freqs = self.forward( | |
| lambda: self.get_seq_pos(seq_len, device=device, dtype=dtype, offset=offset), | |
| cache_key=f'freqs:{seq_len}|offset:{offset}') | |
| return apply_rotary_emb(freqs, t) | |
| def rotate_queries_with_cached_keys(self, q, k, seq_dim=-2): | |
| q_len, k_len = q.shape[seq_dim], k.shape[seq_dim] | |
| assert q_len <= k_len | |
| q = self.rotate_queries_or_keys(q, seq_dim=seq_dim, freq_seq_len=k_len) | |
| k = self.rotate_queries_or_keys(k, seq_dim=seq_dim) | |
| return q, k | |
| def rotate_queries_and_keys(self, q, k, seq_dim=-2): | |
| assert self.use_xpos | |
| device, dtype, seq_len = q.device, q.dtype, q.shape[seq_dim] | |
| seq = self.get_seq_pos(seq_len, dtype=dtype, device=device) | |
| freqs = self.forward(lambda: seq, cache_key=f'freqs:{seq_len}') | |
| scale = self.get_scale(lambda: seq, cache_key=f'scale:{seq_len}').to(dtype) | |
| rotated_q = apply_rotary_emb(freqs, q, scale=scale) | |
| rotated_k = apply_rotary_emb(freqs, k, scale=scale**-1) | |
| return rotated_q, rotated_k | |
| def get_scale(self, t, cache_key=None): | |
| assert self.use_xpos | |
| if exists(cache_key) and cache_key in self.cache: | |
| return self.cache[cache_key] | |
| if callable(t): | |
| t = t() | |
| scale = 1. | |
| if self.use_xpos: | |
| power = (t - len(t) // 2) / self.scale_base | |
| scale = self.scale**rearrange(power, 'n -> n 1') | |
| scale = torch.cat((scale, scale), dim=-1) | |
| if exists(cache_key): | |
| self.cache[cache_key] = scale | |
| return scale | |
| def forward(self, t, cache_key=None): | |
| if exists(cache_key) and cache_key in self.cache: | |
| return self.cache[cache_key] | |
| if callable(t): | |
| t = t() | |
| freqs = self.freqs | |
| freqs = einsum('..., f -> ... f', t.type(freqs.dtype), freqs) | |
| freqs = repeat(freqs, '... n -> ... (n r)', r=2) | |
| if exists(cache_key): | |
| self.cache[cache_key] = freqs | |
| return freqs |