Views
No views yet
transformers 即可运行1from transformers import AutoTokenizer, AutoModelForCausalLM
2
3model_id = "wdndev/tiny_llm_sft_76m_llama"
4
5tokenizer = AutoTokenizer.from_pretrained(model_id)
6model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", trust_remote_code=True)
7
8sys_text = "你是由wdndev开发的个人助手。"
9# user_text = "中国的首都是哪儿?"
10# user_text = "你叫什么名字?"
11user_text = "介绍一下中国"
12input_txt = "\n".join(["<|system|>", sys_text.strip(),
13 "<|user|>", user_text.strip(),
14 "<|assistant|>"]).strip() + "\n"
15
16model_inputs = tokenizer(input_txt, return_tensors="pt").to(model.device)
17generated_ids = model.generate(model_inputs.input_ids, max_new_tokens=200)
18generated_ids = [
19 output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
20]
21response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
22print(response)