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