Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4# 加载模型和tokenizer
5model_name = "cjkasbdkjnlakb/qwen3-4b-agent-0927-qc"
6model = AutoModelForCausalLM.from_pretrained(
7 model_name,
8 torch_dtype=torch.bfloat16,
9 device_map="auto",
10 trust_remote_code=True
11)
12tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
13
14# 推理示例
15def chat(messages):
16 text = tokenizer.apply_chat_template(
17 messages,
18 tokenize=False,
19 add_generation_prompt=True
20 )
21
22 inputs = tokenizer(text, return_tensors="pt").to(model.device)
23
24 with torch.no_grad():
25 outputs = model.generate(
26 **inputs,
27 max_new_tokens=512,
28 do_sample=True,
29 temperature=0.7,
30 top_p=0.9,
31 pad_token_id=tokenizer.eos_token_id
32 )
33
34 response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
35 return response
36
37# 使用示例
38messages = [
39 {"role": "user", "content": "你好,请介绍一下你自己。"}
40]
41response = chat(messages)
42print(response)torch.bfloat16 或 torch.float16 精度以节省显存device_map 设置