File size: 2,797 Bytes
b777314
 
c47d202
 
 
b777314
63f23aa
c5f1c0b
 
 
 
 
 
 
 
 
 
 
 
 
 
5690751
c5f1c0b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
---
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, :].copy()).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
```