1from transformers import AutoTokenizer, AutoModelForCausalLM
2
3# 加载 tokenizer 和模型
4tokenizer = AutoTokenizer.from_pretrained("sds-ai/Yee-270m")
5model = AutoModelForCausalLM.from_pretrained(
6 "sds-ai/Yee-270m",
7 torch_dtype="auto",
8 device_map="auto"
9)
10
11# 输入提示
12prompt = "请帮我检查这份数据是否包含敏感字段?"
13
14# 应用聊天模板
15messages = [{"role": "user", "content": prompt}]
16text = tokenizer.apply_chat_template(
17 messages,
18 tokenize=False,
19 add_generation_prompt=True
20)
21
22# 编码输入
23inputs = tokenizer([text], return_tensors="pt").to(model.device)
24
25# 生成响应
26response_ids = model.generate(**inputs, max_new_tokens=1024)
27response = tokenizer.decode(response_ids[0][len(inputs.input_ids[0]):], skip_special_tokens=True)
28
29print("小熠:\n", response)
1from transformers import pipeline
2
3pipe = pipeline("text-generation", model="sds-ai/Yee-270m")
4response = pipe("数据安全最佳实践有哪些?")