模型参数设计遵循
Qwen2-0.5B，根据测试，深而窄的模型表现比较好，并且对于小模型，FFN 的升维维度可以适当更大。
不同于 SmolLM 在最后 20% 的步骤开始衰减学习率，这里梯形调度的衰减步骤占比达到了 30%，采用和 MiniCPM 一致的指数衰减，最低衰减至最大学习率的 1%。
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3
4model_path = 'Mxode/SmolLM-Chinese-180M'
5
6tokenizer = AutoTokenizer.from_pretrained(model_path)
7model = AutoModelForCausalLM.from_pretrained(model_path).to('cuda:0', torch.float16)
8
9def get_response(text: str, model, **kwargs):
10 generation_args = dict(
11 max_new_tokens = kwargs.pop("max_new_tokens", 512),
12 do_sample = kwargs.pop("do_sample", True),
13 temperature = kwargs.pop("temperature", 0.55),
14 top_p = kwargs.pop("top_p", 0.8),
15 top_k = kwargs.pop("top_k", 40),
16 **kwargs
17 )
18
19 prompt = text
20 model_inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
21
22 generated_ids = model.generate(model_inputs.input_ids, **generation_args)
23 generated_ids = [
24 output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
25 ]
26
27 response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
28 return response
29
30text = "牛奶作为人类日常必须的优良营养食品，"
31
32response = get_response(text, model, max_new_tokens=256, do_sample=True, temperature=1.0)
33print(f'{text}{response}')
