x-transformers 1.42.7__tar.gz → 1.42.8__tar.gz
Sign up to get free protection for your applications and to get access to all the features.
- {x_transformers-1.42.7/x_transformers.egg-info → x_transformers-1.42.8}/PKG-INFO +1 -1
- {x_transformers-1.42.7 → x_transformers-1.42.8}/setup.py +1 -1
- {x_transformers-1.42.7 → x_transformers-1.42.8}/tests/test_x_transformers.py +94 -1
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/x_transformers.py +20 -10
- {x_transformers-1.42.7 → x_transformers-1.42.8/x_transformers.egg-info}/PKG-INFO +1 -1
- {x_transformers-1.42.7 → x_transformers-1.42.8}/LICENSE +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/README.md +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/setup.cfg +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/__init__.py +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/attend.py +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/autoregressive_wrapper.py +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/continuous.py +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/dpo.py +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/multi_input.py +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/neo_mlp.py +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/nonautoregressive_wrapper.py +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/xl_autoregressive_wrapper.py +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/xval.py +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers.egg-info/SOURCES.txt +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers.egg-info/dependency_links.txt +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers.egg-info/requires.txt +0 -0
- {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers.egg-info/top_level.txt +0 -0
@@ -1,12 +1,15 @@
|
|
1
1
|
import pytest
|
2
|
+
|
2
3
|
import torch
|
4
|
+
from torch import nn
|
5
|
+
from torch.nn import Module
|
3
6
|
|
4
7
|
from x_transformers.x_transformers import (
|
5
8
|
XTransformer,
|
6
9
|
TransformerWrapper,
|
7
10
|
Encoder,
|
8
11
|
Decoder,
|
9
|
-
|
12
|
+
LinearNoBias,
|
10
13
|
)
|
11
14
|
|
12
15
|
from x_transformers.neo_mlp import (
|
@@ -394,3 +397,93 @@ def test_custom_alibi():
|
|
394
397
|
pos = torch.tensor([[0, 1, 2, 4], [1, 3, 5, 7]])
|
395
398
|
|
396
399
|
logits = model(x, pos = pos)
|
400
|
+
|
401
|
+
|
402
|
+
@pytest.mark.parametrize('embedder_type', ('embedding', 'none', 'custom'))
|
403
|
+
def test_embedder(embedder_type):
|
404
|
+
num_tokens = 20000
|
405
|
+
dim = 128
|
406
|
+
token_emb_kwargs = {}
|
407
|
+
|
408
|
+
if embedder_type == 'embedding':
|
409
|
+
embedder = nn.Embedding(num_tokens, dim)
|
410
|
+
elif embedder_type == 'none':
|
411
|
+
embedder = None
|
412
|
+
else:
|
413
|
+
class CustomEmbedder(Module):
|
414
|
+
"""
|
415
|
+
Made up embedder that sums two embeddings. Just to check if we can pass additional input to the embedder's
|
416
|
+
forward pass without breaking the model.
|
417
|
+
"""
|
418
|
+
def __init__(self, num_tokens, dim):
|
419
|
+
super().__init__()
|
420
|
+
self.embed_x = nn.Embedding(num_tokens, dim)
|
421
|
+
self.embed_y = nn.Embedding(num_tokens, dim)
|
422
|
+
|
423
|
+
def forward(self, x, y):
|
424
|
+
return self.embed_x(x) + self.embed_y(y)
|
425
|
+
|
426
|
+
def init_(self):
|
427
|
+
pass
|
428
|
+
|
429
|
+
embedder = CustomEmbedder(num_tokens, dim)
|
430
|
+
token_emb_kwargs['y'] = torch.randint(0, num_tokens, (2, 1024))
|
431
|
+
|
432
|
+
model = TransformerWrapper(
|
433
|
+
num_tokens = num_tokens,
|
434
|
+
max_seq_len = 1024,
|
435
|
+
attn_layers = Decoder(
|
436
|
+
dim = dim,
|
437
|
+
depth = 6,
|
438
|
+
heads = 8,
|
439
|
+
),
|
440
|
+
token_emb = embedder,
|
441
|
+
)
|
442
|
+
|
443
|
+
x = torch.randint(0, 20000, (2, 1024))
|
444
|
+
|
445
|
+
output = model(x, token_emb_kwargs=token_emb_kwargs)
|
446
|
+
assert output.shape == (2, 1024, 20000)
|
447
|
+
|
448
|
+
|
449
|
+
@pytest.mark.parametrize("to_logits", ('linear', 'none', 'pointer'))
|
450
|
+
def test_to_logits(to_logits):
|
451
|
+
num_tokens = 20000
|
452
|
+
dim = 128
|
453
|
+
|
454
|
+
to_logits_kwargs = {}
|
455
|
+
|
456
|
+
if to_logits == 'linear':
|
457
|
+
logit_mapper = LinearNoBias(dim, num_tokens)
|
458
|
+
elif to_logits == 'none':
|
459
|
+
logit_mapper = None
|
460
|
+
else:
|
461
|
+
class PointerNetworkLogits(Module):
|
462
|
+
def __init__(self, dim):
|
463
|
+
super().__init__()
|
464
|
+
self.proj_to_pointers = nn.Linear(dim, dim)
|
465
|
+
|
466
|
+
def forward(self, model_embeddings, input_embeddings):
|
467
|
+
pointers = self.proj_to_pointers(model_embeddings)
|
468
|
+
logits = torch.matmul(pointers, input_embeddings.permute(0, 2, 1))
|
469
|
+
return logits
|
470
|
+
|
471
|
+
logit_mapper = PointerNetworkLogits(dim)
|
472
|
+
to_logits_kwargs['input_embeddings'] = torch.randn(2, 20000, dim)
|
473
|
+
|
474
|
+
model = TransformerWrapper(
|
475
|
+
num_tokens = num_tokens,
|
476
|
+
max_seq_len = 1024,
|
477
|
+
attn_layers = Decoder(
|
478
|
+
dim = dim,
|
479
|
+
depth = 6,
|
480
|
+
heads = 8,
|
481
|
+
),
|
482
|
+
to_logits = logit_mapper,
|
483
|
+
)
|
484
|
+
|
485
|
+
x = torch.randint(0, num_tokens, (2, 1024))
|
486
|
+
|
487
|
+
output = model(x, to_logits_kwargs=to_logits_kwargs)
|
488
|
+
|
489
|
+
assert output.shape == (2, 1024, 20000)
|
@@ -238,6 +238,13 @@ class TokenEmbedding(Module):
|
|
238
238
|
token_emb = self.emb(x.long())
|
239
239
|
return l2norm(token_emb) if self.l2norm_embed else token_emb
|
240
240
|
|
241
|
+
def init_(self):
|
242
|
+
if self.l2norm_embed:
|
243
|
+
nn.init.normal_(self.emb.weight, std=1e-5)
|
244
|
+
return
|
245
|
+
nn.init.kaiming_normal_(self.emb.weight)
|
246
|
+
|
247
|
+
|
241
248
|
# positional embeddings
|
242
249
|
|
243
250
|
class AbsolutePositionalEmbedding(Module):
|
@@ -2261,7 +2268,8 @@ class TransformerWrapper(Module):
|
|
2261
2268
|
token_emb: TokenEmbedding | None = None,
|
2262
2269
|
mixture_of_softmax = False,
|
2263
2270
|
mixture_of_softmax_k = 4,
|
2264
|
-
sigsoftmax_logits = False
|
2271
|
+
sigsoftmax_logits = False,
|
2272
|
+
to_logits: Module | None = None,
|
2265
2273
|
):
|
2266
2274
|
super().__init__()
|
2267
2275
|
|
@@ -2363,11 +2371,12 @@ class TransformerWrapper(Module):
|
|
2363
2371
|
if return_only_embed:
|
2364
2372
|
self.to_logits = None
|
2365
2373
|
elif tie_embedding:
|
2374
|
+
assert isinstance(token_emb, TokenEmbedding), 'can only tie embedding if using `TokenEmbedding`'
|
2366
2375
|
self.to_logits = lambda t: t @ self.token_emb.emb.weight.t()
|
2367
2376
|
elif num_output_heads > 1:
|
2368
2377
|
self.to_logits = ModuleList([LinearNoBias(dim, logits_dim) for _ in range(num_output_heads)])
|
2369
2378
|
else:
|
2370
|
-
self.to_logits = LinearNoBias(dim, logits_dim)
|
2379
|
+
self.to_logits = LinearNoBias(dim, logits_dim) if not exists(to_logits) else to_logits
|
2371
2380
|
|
2372
2381
|
# memory tokens (like [cls]) from Memory Transformers paper
|
2373
2382
|
|
@@ -2388,13 +2397,12 @@ class TransformerWrapper(Module):
|
|
2388
2397
|
self.can_cache_kv_outside_max_seq_len = no_abs_pos_emb
|
2389
2398
|
|
2390
2399
|
def init_(self):
|
2400
|
+
if hasattr(self.token_emb, 'init_'):
|
2401
|
+
self.token_emb.init_()
|
2402
|
+
|
2391
2403
|
if self.l2norm_embed:
|
2392
|
-
nn.init.normal_(self.token_emb.emb.weight, std = 1e-5)
|
2393
2404
|
if not isinstance(self.pos_emb, always):
|
2394
2405
|
nn.init.normal_(self.pos_emb.emb.weight, std = 1e-5)
|
2395
|
-
return
|
2396
|
-
|
2397
|
-
nn.init.kaiming_normal_(self.token_emb.emb.weight)
|
2398
2406
|
|
2399
2407
|
def forward(
|
2400
2408
|
self,
|
@@ -2417,7 +2425,9 @@ class TransformerWrapper(Module):
|
|
2417
2425
|
attn_z_loss_weight = 1e-4,
|
2418
2426
|
seq_start_pos = None,
|
2419
2427
|
cache: LayerIntermediates | None = None,
|
2420
|
-
|
2428
|
+
token_emb_kwargs = dict(),
|
2429
|
+
to_logits_kwargs = dict(),
|
2430
|
+
**kwargs,
|
2421
2431
|
):
|
2422
2432
|
b, n, device, num_mems, has_memory_tokens, emb_frac_gradient, orig_mask = x.shape[0], x.shape[1], x.device, self.num_memory_tokens, self.num_memory_tokens > 0, self.emb_frac_gradient, mask
|
2423
2433
|
|
@@ -2428,7 +2438,7 @@ class TransformerWrapper(Module):
|
|
2428
2438
|
|
2429
2439
|
external_pos_emb = exists(pos) and pos.dtype != torch.long
|
2430
2440
|
pos_emb = self.pos_emb(x, pos = pos, seq_start_pos = seq_start_pos) if not external_pos_emb else pos
|
2431
|
-
x = self.token_emb(x) + pos_emb
|
2441
|
+
x = self.token_emb(x, **token_emb_kwargs) + pos_emb
|
2432
2442
|
|
2433
2443
|
# add additional embeddings
|
2434
2444
|
|
@@ -2583,9 +2593,9 @@ class TransformerWrapper(Module):
|
|
2583
2593
|
|
2584
2594
|
if not return_embeddings:
|
2585
2595
|
if self.has_multiple_heads:
|
2586
|
-
logits = tuple(fn(x) for fn in self.to_logits)
|
2596
|
+
logits = tuple(fn(x, **to_logits_kwargs) for fn in self.to_logits)
|
2587
2597
|
else:
|
2588
|
-
logits = self.to_logits(x)
|
2598
|
+
logits = self.to_logits(x, **to_logits_kwargs)
|
2589
2599
|
|
2590
2600
|
# maybe sig softmax
|
2591
2601
|
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
{x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers.egg-info/dependency_links.txt
RENAMED
File without changes
|
File without changes
|
File without changes
|