Update README.md

README.md

@@ -2,4 +2,99 @@
 license: apache-2.0
 ---
 
-Train in 30B Byte. Mode size 353M. Table 2 in [MambaByte](https://arxiv.org/abs/2401.13660)
+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
+
+    def _initialize(self):
+        """idx is the channel of the inpu
+```