|
--- |
|
license: apache-2.0 |
|
metrics: |
|
- perplexity |
|
pipeline_tag: text-generation |
|
--- |
|
|
|
Train in 30B Byte. Mode size 353M. Table 2 in [MambaByte](https://arxiv.org/abs/2401.13660) |
|
|
|
To use |
|
|
|
``` |
|
import torch |
|
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel |
|
|
|
import numpy as np |
|
|
|
model=MambaLMHeadModel.from_pretrained("JunxiongWang/MambaByte_Code", device='cuda', dtype=torch.float32) |
|
|
|
text = "import torch" |
|
text_byte = np.frombuffer(text.encode('utf-8'), dtype=np.uint8) |
|
input_ids = torch.from_numpy(text_byte[None, :]).long().cuda() |
|
|
|
sample = model.generate( |
|
input_ids=input_ids, |
|
max_length=2048, |
|
cg=True, |
|
return_dict_in_generate=True, |
|
output_scores=True, |
|
enable_timing=True, |
|
temperature=1, |
|
top_k=256, |
|
top_p=0.9, |
|
) |
|
|
|
print(bytes(sample.sequences[0].tolist()).decode('utf-8')) |
|
``` |
|
|
|
Output |
|
|
|
``` |
|
import torch |
|
import numpy as np |
|
import torch.nn.functional as F |
|
from torch.autograd import Variable |
|
|
|
from networkx.states import TransientState |
|
|
|
def extract_data(num_epochs, epochs, is_last_epoch): |
|
|
|
def get_data(num_features, num_classes): |
|
data_features = num_features |
|
data_classes = num_classes |
|
data_labels = num_epochs |
|
|
|
if num_features == 0 or num_classes == 0: |
|
return data_features, data_classes |
|
if is_last_epoch: |
|
data_features = num_features |
|
data_classes = num_classes |
|
data_labels = num_epochs |
|
return data_features, data_classes |
|
|
|
data_features, data_classes = get_data(num_epochs, epochs, is_last_epoch) |
|
data_labels = num_epochs * 2 |
|
return data_features, data_classes |
|
|
|
|
|
class NumChannel: |
|
def __init__(self, x, y, dx=1, dy=1, idx=1, data_size=2, epoch=None): |
|
"""idx is the channel index with data feature in the first epoch. |
|
x is the channel of the input data. |
|
y is the element of the input data. |
|
dx is the element of the data feature of the input data. |
|
data_size is the size of the element of the data. |
|
epoch is the channel of the element of the data. |
|
""" |
|
self.x = x |
|
self.y = y |
|
self.dx = dx |
|
self.data_size = data_size |
|
self.epoch = epoch |
|
self.reference_count = 0 |
|
self.data_features = {} |
|
self.data_classes = {} |
|
|
|
self._initialize() |
|
if idx is not None: |
|
self._start_time = time.time() |
|
|
|
def _initialize(self): |
|
"""idx is the channel index with data feature in the first epoch. |
|
x is the channel of the input data. |
|
y is the element of the input data. |
|
dx is the element of the data feature of the input data. |
|
data_size is the size of the element of the data. |
|
epoch is the channel of the element of the data. |
|
""" |
|
self.idx = idx |
|
``` |
