glm-moe-dsa-tiny-random

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "yujiepan/glm-moe-dsa-tiny-random"
tokenizer = AutoTokenizer.from_pretrained(model_id)
input_ids = torch.randint(1000, 2000, size=(1, 2333), dtype=torch.long).cuda()  # trigger DSA
# messages = [{"role": "user", "content": "hello"}]
# input_ids = tokenizer(messages, return_tensors="pt").input_ids.cuda()
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype=torch.bfloat16,
    device_map="cuda",
)
generated_ids = model.generate(input_ids, max_new_tokens=32)
output_text = tokenizer.decode(generated_ids[0][input_ids.shape[1]:])
print(output_text)