Views
No views yet
1import transformers
2from peft import PeftModel, PeftConfig
3from transformers import AutoModelForCausalLM, AutoTokenizer
4import torch
5from torch import cuda, bfloat16
6
7base_model_id = 'meta-llama/Llama-2-7b-chat-hf'
8
9device = f'cuda:{cuda.current_device()}' if cuda.is_available() else 'cpu'
10
11bnb_config = transformers.BitsAndBytesConfig(
12 load_in_4bit=True,
13 bnb_4bit_quant_type='nf4',
14 bnb_4bit_use_double_quant=True,
15 bnb_4bit_compute_dtype=bfloat16
16)
17
18
19hf_auth = "your-hf-access-token"
20model_config = transformers.AutoConfig.from_pretrained(
21 base_model_id,
22 use_auth_token=hf_auth
23)
24
25model = transformers.AutoModelForCausalLM.from_pretrained(
26 base_model_id,
27 trust_remote_code=True,
28 config=model_config,
29 quantization_config=bnb_config,
30 device_map='auto',
31 use_auth_token=hf_auth
32)
33
34config = PeftConfig.from_pretrained("Ashishkr/llama2-call-summarization")
35model = PeftModel.from_pretrained(model, "Ashishkr/llama2-call-summarization").to(device)
36
37model.eval()
38print(f"Model loaded on {device}")
39
40tokenizer = transformers.AutoTokenizer.from_pretrained(
41 base_model_id,
42 use_auth_token=hf_auth
43)
44
451
2
3def llama_generate(
4 model: AutoModelForCausalLM,
5 tokenizer: AutoTokenizer,
6 prompt: str,
7 max_new_tokens: int = 128,
8 temperature: float = 0.92):
9
10 inputs = tokenizer(
11 [prompt],
12 return_tensors="pt",
13 return_token_type_ids=False,
14 ).to(
15 device
16 )
17
18 # Check if bfloat16 is supported, otherwise use float16
19 dtype_to_use = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
20
21 with torch.autocast("cuda", dtype=dtype_to_use):
22 response = model.generate(
23 **inputs,
24 max_new_tokens=max_new_tokens,
25 temperature=temperature,
26 return_dict_in_generate=True,
27 eos_token_id=tokenizer.eos_token_id,
28 pad_token_id=tokenizer.pad_token_id,
29 )
30
31 decoded_output = tokenizer.decode(
32 response["sequences"][0],
33 skip_special_tokens=True,
34 )
35
36 return decoded_output[len(prompt) :]
37
38prompt = """
39 instruction: "summarize this conversation :" \n
40
41input: "Oli: I've talked to some people from the third year
42Jacob: About the statistics exam?
43Marcia: What did they say?
44Oli: Yeah, about the exam Oli: We need to prepare for a battle
45Jacob: So it will be difficult
46Oli: They said it was the hardest exam ever
47Marcia: 😱
48Oli: The questions were displayed on the screen
49Oli: One minute per question and it disappears
50Oli: They won't come back so if you didn't get your answer you're fucked
51Marcia: So we need to make the calculations really fast
52Jacob: That's insane
53Oli: I know
54Oli: Very stressful
55Marcia: How are we even supposed to study for it?
56Marcia: With a timer?
57Oli: I guess
58Marcia: Did anybody pass it last year
59Oli: Some people did, but the majority had to take the second or even the third chance"\n
60
61response: """
62response = llama_generate(
63 model,
64 tokenizer,
65 prompt,
66 max_new_tokens=100,
67 temperature=0.9,
68).split("<eos>")[0].strip()
69
70print(response)