JunxiongWang/MambaByte_Stories

Train in 30B Byte. Mode size 353M. Table 2 in MambaByte

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, :].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

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 .