Views
No views yet
1import transformers
2from transformers import AutoTokenizer, AutoModelForCausalLM
3import torch
4
5model_id = "UnicomAI/Unichat-llama3-Chinese-8B-28K"
6
7pipeline = transformers.pipeline(
8 "text-generation",
9 model=model_id,
10 model_kwargs={"torch_dtype": torch.bfloat16},
11 device="cuda",
12)
13
14
15messages = [
16 {"role": "system", "content": "A chat between a curious user and an artificial intelligence assistant.The assistant gives helpful, detailed, and polite answers to the user's questions."},
17 {"role": "user", "content": "你是谁"},
18]
19
20
21prompt = pipeline.tokenizer.apply_chat_template(
22 messages,
23 tokenize=False,
24 add_generation_prompt=True
25)
26
27terminators = [
28 pipeline.tokenizer.eos_token_id,
29 pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")
30]
31
32
33outputs = pipeline(
34 prompt,
35 max_new_tokens=2048,
36 eos_token_id=terminators,
37 do_sample=True,
38 temperature=0.6,
39 top_p=1,
40 repetition_penalty=1.15
41)
42print(outputs[0]["generated_text"][len(prompt):])