Update README.md

README.md

@@ -2,4 +2,59 @@
 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_Stories", device='cuda', dtype=torch.float32)
+
+text = "It was terribly cold and nearly dark on the last evening of the old year, and the snow was falling fast."
+
+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
+
+```
+It was terribly cold and nearly dark on the last evening of the old year, and the snow was falling fast.
+It had been a warm summer , but the ground was dry and cracked .
+The grass was like a great cloud , as deep as a city of ancient people .
+Strange to see a land with so few colors standing strong on the land ; but there it was , passing by in her dreams , the colors of flowers and grasses and birds .
+Of course , the old year was all a faded charcoal black , somewhat like the distant gloaming in the city where she lived .
+She seemed to be walking along the shore of a land that was no longer there .
+And indeed , at first it seemed as if there was another world there , a world inhabited by beings no one had ever seen before , who all consisted of flowers and small fish .
+And they were lovely people , but of whom had never the least idea of a love of anything or anyone before .
+They were neither fierce nor merciless , but deadly serious and bloodthirsty .
+Their deaths were against nature , and nature alone knew how to play tricks on those who opposed it .
+The living , the lived and dead , were all the same to the living ; the living ones were no more than living things whose deaths were inconceivable to them .
+Who could understand why the living were plagued by these men and women ?
+But at least the living were dead ; no doubt they would eat themselves to death in the end .
+The sun shone and the warmth of the day renewed its purity .
+The living were finally able to enjoy the refreshing waters that often flow from the spring .
+The sun shone brightly , and the air was filled with the scent of the flowers .
+But the flowers had no soul .
+They were cold and lonely , and would rather have died in the cold and the wet , than be with the living ones who were dead .
+Of course , this was no great loss , for the living could not take a death in the cold or the weather , for the living could only die by coldness .
+They were cold because they knew that no human being could take warmth and shelter from the cold of death .
+```