关键代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("XiaomiMiMo/MiMo-7B-RL")
model = AutoModelForCausalLM.from_pretrained("XiaomiMiMo/MiMo-7B-RL", device_map="auto", torch_dtype="auto")
messages = [
    {"role": "user", "content": "如果晾干5件衬衫需要4小时，那么晾干20件衬衫需要多久？"},
]
inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")

outputs = model.generate(
    inputs.to(model.device),
    max_new_tokens=4096,
    temperature=0.8,
    top_p=0.95,
    do_sample=True,
)
print(tokenizer.decode(outputs[0]))