Views
No views yet
1import argparse
2from transformers import AutoTokenizer, pipeline
3
4# Attention! The directory must not contain "." and can be replaced with "_".
5parser = argparse.ArgumentParser()
6parser.add_argument('--model_path', default="IndexTeam/Index-1.9B-Chat", type=str, help="")
7parser.add_argument('--device', default="cpu", type=str, help="") # also could be "cuda" or "mps" for Apple silicon
8args = parser.parse_args()
9
10tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
11generator = pipeline("text-generation",
12 model=args.model_path,
13 tokenizer=tokenizer, trust_remote_code=True,
14 device=args.device)
15
16
17system_message = "你是由哔哩哔哩自主研发的大语言模型,名为“Index”。你能够根据用户传入的信息,帮助用户完成指定的任务,并生成恰当的、符合要求的回复。"
18query = "续写 天不生我金坷垃"
19model_input = []
20model_input.append({"role": "system", "content": system_message})
21model_input.append({"role": "user", "content": query})
22
23model_output = generator(model_input, max_new_tokens=300, top_k=5, top_p=0.8, temperature=0.3, repetition_penalty=1.1, do_sample=True)
24
25print('User:', query)
26print('Model:', model_output)