use threshold of ema cluster sizes to determine which codes to replac…

…e, as in soundstream paper, thanks to @wesbz

README.md

@@ -110,7 +110,7 @@ quantized, indices, commit_loss = vq(x)
 
 ### Expiring stale codes
 
-Finally, the SoundStream paper has a scheme where they replace codes that have not been used in a certain number of consecutive batches with a randomly selected vector from the current batch. You can set this threshold for consecutive misses before replacement with `max_codebook_misses_before_expiry` keyword. (I know it is a bit long, but I couldn't think of a better name)
+Finally, the SoundStream paper has a scheme where they replace codes that have hits below a certain threshold with randomly selected vector from the current batch. You can set this threshold with `threshold_ema_dead_code` keyword.
 
 ```python
 import torch
@@ -119,7 +119,7 @@ from vector_quantize_pytorch import VectorQuantize
 vq = VectorQuantize(
     dim = 256,
     codebook_size = 512,
-    max_codebook_misses_before_expiry = 5  # should actively replace any codes that were missed 5 times in a row during training
+    threshold_ema_dead_code = 2  # should actively replace any codes that have an exponential moving average cluster size less than 2
 )
 
 x = torch.randn(1, 1024, 256)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '0.3.3',
+  version = '0.3.4',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   author = 'Phil Wang',