Update README.md

README.md

@@ -2,4 +2,47 @@
 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
+import numpy as np
+
+from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
+
+model=MambaLMHeadModel.from_pretrained("JunxiongWang/MambaByte_Arxiv", device='cuda', dtype=torch.bfloat16)
+
+text = "\documentclass[12pt]{article}"
+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:
+
+```
+\documentclass[12pt]{article}}}}^{{\mathbf{P}}\uplus{\mathbf{Q}}}}}}}{}}$ is a symmetric poset. This implies that $$\operatorname{end}({\mathscr{L}}) = \operatorname{end}({\mathscr{L}}\setminus\{\sigma_{{\mathbf{P}}}\}) = \operatorname{end}({\mathscr{L}}\setminus\{\sigma_{{\mathbf{Q}}}\}) = \operatorname{end}({\mathscr{L}}\setminus\{\sigma_{{\mathbf{P}}},\sigma_{{\mathbf{Q}}}\}),$$ i.e., ${\mathscr{L}}$ is $\{\sigma_{{\mathbf{P}}},\sigma_{{\mathbf{Q}}}\}$-bistochastic for any ${\mathbf{P}}\neq{\mathbf{Q}}$. Thus, ${\mathscr{L}}$ is reversible, and is in fact maximal among all $\{\sigma_{{\mathbf{P}}},\sigma_{{\mathbf{Q}}}\}$-bistochastic matrices.
+
+Since ${\mathscr{L}}$ is in the same class as $\sigma_{{\mathbf{P}}}^{{\mathbf{Q}}}$, we have $\operatorname{end}({\mathscr{L}})\subseteq\operatorname{end}({\mathscr{L}})$. Conversely, if $\operatorname{end}({\mathscr{L}})=\operatorname{end}({\mathscr{L}})$, then $\sigma_{{\mathbf{P}}}^{{\mathbf{Q}}}$ is maximal in $\operatorname{end}({\mathscr{L}})$. Since ${\mathbf{P}}\setminus\{\sigma_{{\mathbf{P}}}\}\subseteq\operatorname{end}({\mathscr{L}})$, this implies that ${\mathscr{L}}$ is in the same class as $\sigma_{{\mathbf{P}}}^{{\mathbf{Q}}}$, and hence ${\mathscr{L}}$ is reversible.
+
+We are now ready to show that $\{\sigma_{{\mathbf{P}}},\sigma_{{\mathbf{Q}}}\}$-bistochastic matrices form a symmetric poset of ends.
+
+\[lem:end\_symm\_class\] Let ${\mathbf{P}},{\mathbf{Q}}\in{\mathscr{M}}$. Then $\sigma_{{\mathbf{P}}}^{{\mathbf{Q}}}$ is symmetric if and only if $\operatorname{end}({\mathscr{L}})=\operatorname{end}({\mathscr{L}})$.
+
+Suppose that $\operatorname{end}({\mathscr{L}})=\operatorname{end}({\mathscr{L}})$, and we prove that $\sigma_{{\mathbf{P}}}^{{\mathbf{Q}}}$ is symmetric. Clearly, $\operatorname{end}({\mathscr{L}})$ contains exactly the ends of $\operatorname{end}({\mathscr{L}})$ by definition, and the only case that survives is when $\operatorname{end}({\mathscr{L}})=\operatorname{end}({\mathscr{L}})$. By construction, this means that $\sigma_{{\mathbf{P}}}
+```