Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
2import torch
3
4
5if __name__ == '__main__':
6 base_model = "meta-llama/Meta-Llama-3-8B-Instruct"
7 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
8
9 bnb_config = BitsAndBytesConfig(
10 load_in_4bit=True,
11 bnb_4bit_quant_type="nf4",
12 bnb_4bit_compute_dtype=torch.float16,
13 bnb_4bit_use_double_quant=True,
14 )
15
16 model = AutoModelForCausalLM.from_pretrained(
17 base_model,
18 quantization_config=bnb_config,
19 device_map={"": 0},
20 attn_implementation="eager"
21 )
22
23 tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
24
25 tokenizer.pad_token = '<|pad|>'
26 tokenizer.pad_token_id = 128255
27
28 #Load LORA weights
29 model.load_adapter("Anonymous-pre-publication/FoodSEM-LLM")
30 model.config.use_cache = True
31 model.eval()
32
33 system_prompt = ""
34 user_prompt = "Please, may we have links to the Hansard taxonomy for these entities provided: soft butter, mango, daiquiri mixer, maple extract, salt, anise flavored liqueur, hemp seeds, yeast mixture, thighs?"
35
36 messages = [
37 {
38 "role": "user",
39 "content": f"{system_prompt} {user_prompt}".strip()
40 }
41 ]
42
43 prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
44
45 #Here we have a batch of one
46 tokenizer_input = [prompt]
47
48 inputs = tokenizer(tokenizer_input, return_tensors="pt", padding=True, truncation=True, max_length=1024).to(device)
49 generated_ids = model.generate(**inputs, max_new_tokens=1024, do_sample=True)
50 answers = tokenizer.batch_decode(generated_ids[:, inputs['input_ids'].shape[1]:])
51 answers = [x.split('<|eot_id|>')[0].strip() for x in answers]
52 print(answers)