longnet class

LongNet/model.py

@@ -71,21 +71,28 @@ def forward(self, text_tokens, **kwargs):
  return self.decoder(model_input, passed_x=model_input)[0]
 
 
-# class LongNet(Module):
-# def __init__(self):
-# super().__init__()
-
-# self.model = LongNet(
-# num_tokens = 16000, # number of tokens
-# dim = (512, 256), # transformer model dimension (512 for coarsest, 256 for fine in this example)
-# max_seq_len = (1024, 4), # sequence length for global and then local. this can be more than 2
-# depth = (6, 4), # number of layers for global and then local. this can be more than 2, but length must match the max_seq_len's
-# dim_head = 64, # dimension per head
-# heads = 8, # number of attention heads
-# flash_attn = True # use flash attention
-# )
-
-# def forward(self, text_tokens, temperature: int = None, filter_thres: int = None, **kwargs):
-# sampled = self.model.generate(temperature=temperature, filter_thres=filter_thres)
-# return sampled
+class DilatedLongNet(Module):
+ def __init__(self):
+ super().__init__()
+
+ self.model = LongNet(
+ num_tokens = 16000, # number of tokens
+ dim = (512, 256), # transformer model dimension (512 for coarsest, 256 for fine in this example)
+ max_seq_len = (1024, 4), # sequence length for global and then local. this can be more than 2
+ depth = (6, 4), # number of layers for global and then local. this can be more than 2, but length must match the max_seq_len's
+ dim_head = 64, # dimension per head
+ heads = 8, # number of attention heads
+ flash_attn = True, # use flash attention
+ dilation_rate = 1, # dilation rate for DilatedAttention
+ segment_size = 0, # segment size for DilatedAttention
+ casual = False, # whether to use causal attention for DilatedAttention
+ use_xpos = False, # whether to use absolute positional embeddings for DilatedAttention
+ use_rel_pos_bias = False, # whether to use relative positional bias for DilatedAttention
+ distributed = False # whether to distribute attention for DilatedAttention
+ )
+
+ def forward(self, text_tokens, temperature: int = None, filter_thres: int = None, **kwargs):
+ sampled = self.model.generate(temperature=temperature, filter_thres=filter_thres)
+ return sampled
+