Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4model_name_or_path = 'YeungNLP/firefly-mixtral-8x7b'
5max_new_tokens = 500
6top_p = 0.9
7temperature = 0.35
8repetition_penalty = 1.0
9
10model = AutoModelForCausalLM.from_pretrained(
11 model_name_or_path,
12 trust_remote_code=True,
13 low_cpu_mem_usage=True,
14 torch_dtype=torch.float16,
15 device_map='auto'
16)
17model = model.eval()
18tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
19
20text = "Compose an engaging travel blog post about a recent trip to Hawaii, highlighting cultural experiences and must-see attractions."
21
22inst_begin_tokens = tokenizer.encode('[INST]', add_special_tokens=False)
23inst_end_tokens = tokenizer.encode('[/INST]', add_special_tokens=False)
24human_tokens = tokenizer.encode(text, add_special_tokens=False)
25input_ids = [tokenizer.bos_token_id] + inst_begin_tokens + human_tokens + inst_end_tokens
26
27# input_ids = human_tokens
28input_ids = torch.tensor([input_ids], dtype=torch.long).cuda()
29
30with torch.no_grad():
31 outputs = model.generate(
32 input_ids=input_ids, max_new_tokens=max_new_tokens, do_sample=True,
33 top_p=top_p, temperature=temperature, repetition_penalty=repetition_penalty,
34 eos_token_id=tokenizer.eos_token_id
35 )
36outputs = outputs.tolist()[0][len(input_ids[0]):]
37response = tokenizer.decode(outputs)
38response = response.strip().replace(tokenizer.eos_token, "").strip()
39print("Chatbot：{}".format(response))
40