国产算力新体验：在沐曦曦云C500上部署小模型的实践

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import os

# 1. 环境配置与模型加载
# 指定本地模型路径，避免重复下载
model_name = "/mnt/moark-models/Qwen3-8B"

# 加载Tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 加载模型
# device_map="auto" 依赖 accelerate 库自动分配显存
# torch_dtype 指定精度以节省显存并提升吞吐
model = AutoModelForCausalLM.from_pretrained(
      model_name,
      torch_dtype=torch.bfloat16, 
      device_map="auto"
)

# 2. 推理逻辑
prompt = "生成一段短文小说，题材不限，但需要在100个字符以上"
messages = [{"role": "user", "content": prompt}]

# 应用Chat Template
text = tokenizer.apply_chat_template(
      messages,
      tokenize=False,
      add_generation_prompt=True,
      enable_thinking=True # 需高版本transformers支持
)

model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# 生成
generated_ids = model.generate(
      **model_inputs,
      max_new_tokens=32768
)

# 3. 结果解析与持久化
output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()

# 解析Thinking部分与Response部分
# 假设151668为结束符token id
try:
      index = len(output_ids) - output_ids[::-1].index(151668)
except ValueError:
      index = 0

thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")

print("=== Thinking ===")
print(thinking_content)
print("\n=== Response ===")
print(content)

# 4. 结果保存
save_path = "/data/exam/text_inference.txt"
os.makedirs(os.path.dirname(save_path), exist_ok=True)
with open(save_path, "w", encoding="utf-8") as f:
    f.write(content)
print(f"\n✅ 结果已保存至: {save_path}")