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.
Files changed (22) hide show
  1. {x_transformers-1.42.7/x_transformers.egg-info → x_transformers-1.42.8}/PKG-INFO +1 -1
  2. {x_transformers-1.42.7 → x_transformers-1.42.8}/setup.py +1 -1
  3. {x_transformers-1.42.7 → x_transformers-1.42.8}/tests/test_x_transformers.py +94 -1
  4. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/x_transformers.py +20 -10
  5. {x_transformers-1.42.7 → x_transformers-1.42.8/x_transformers.egg-info}/PKG-INFO +1 -1
  6. {x_transformers-1.42.7 → x_transformers-1.42.8}/LICENSE +0 -0
  7. {x_transformers-1.42.7 → x_transformers-1.42.8}/README.md +0 -0
  8. {x_transformers-1.42.7 → x_transformers-1.42.8}/setup.cfg +0 -0
  9. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/__init__.py +0 -0
  10. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/attend.py +0 -0
  11. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/autoregressive_wrapper.py +0 -0
  12. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/continuous.py +0 -0
  13. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/dpo.py +0 -0
  14. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/multi_input.py +0 -0
  15. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/neo_mlp.py +0 -0
  16. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/nonautoregressive_wrapper.py +0 -0
  17. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/xl_autoregressive_wrapper.py +0 -0
  18. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers/xval.py +0 -0
  19. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers.egg-info/SOURCES.txt +0 -0
  20. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers.egg-info/dependency_links.txt +0 -0
  21. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers.egg-info/requires.txt +0 -0
  22. {x_transformers-1.42.7 → x_transformers-1.42.8}/x_transformers.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: x-transformers
3
- Version: 1.42.7
3
+ Version: 1.42.8
4
4
  Summary: X-Transformers - Pytorch
5
5
  Home-page: https://github.com/lucidrains/x-transformers
6
6
  Author: Phil Wang
@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
3
3
  setup(
4
4
  name = 'x-transformers',
5
5
  packages = find_packages(exclude=['examples']),
6
- version = '1.42.7',
6
+ version = '1.42.8',
7
7
  license='MIT',
8
8
  description = 'X-Transformers - Pytorch',
9
9
  author = 'Phil Wang',
@@ -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
- AutoregressiveWrapper,
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
- **kwargs
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
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: x-transformers
3
- Version: 1.42.7
3
+ Version: 1.42.8
4
4
  Summary: X-Transformers - Pytorch
5
5
  Home-page: https://github.com/lucidrains/x-transformers
6
6
  Author: Phil Wang
File without changes