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-huggingface-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/llama-2-medical-consultation")
35model = PeftModel.from_pretrained(model, "Ashishkr/llama-2-medical-consultation").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
45
461def llama_generate(
2 model: AutoModelForCausalLM,
3 tokenizer: AutoTokenizer,
4 prompt: str,
5 max_new_tokens: int = 128,
6 temperature: float = 0.92):
7
8 inputs = tokenizer(
9 [prompt],
10 return_tensors="pt",
11 return_token_type_ids=False,
12 ).to(
13 device
14 )
15
16 # Check if bfloat16 is supported, otherwise use float16
17 dtype_to_use = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
18
19 with torch.autocast("cuda", dtype=dtype_to_use):
20 response = model.generate(
21 **inputs,
22 max_new_tokens=max_new_tokens,
23 temperature=temperature,
24 return_dict_in_generate=True,
25 eos_token_id=tokenizer.eos_token_id,
26 pad_token_id=tokenizer.pad_token_id,
27 )
28
29 decoded_output = tokenizer.decode(
30 response["sequences"][0],
31 skip_special_tokens=True,
32 )
33
34 return decoded_output[len(prompt) :]
35
36prompt = """
37 instruction: "If you are a doctor, please answer the medical questions based on the patient's description." \n
38
39input: "Hi, I had a subarachnoid bleed and coiling of brain aneurysm last year.
40I am having some major bilateral temple pain along with numbness that comes and
41goes in my left arm/hand/fingers. I have had headaches since the aneurysm,
42but this is different. Also, my moods have been horrible for the past few weeks.\n
43
44response: """
45# You can use the function as before
46response = llama_generate(
47 model,
48 tokenizer,
49 prompt,
50 max_new_tokens=100,
51 temperature=0.92,
52)
53
54print(response)
55
561
2
3import torch
4import transformers
5from torch import cuda, bfloat16
6from peft import PeftModel, PeftConfig
7from transformers import AutoModelForCausalLM, AutoTokenizer
8
9
10base_model_id = 'meta-llama/Llama-2-7b-chat-hf'
11
12device = f'cuda:{cuda.current_device()}' if cuda.is_available() else 'cpu'
13
14bnb_config = transformers.BitsAndBytesConfig(
15 llm_int8_enable_fp32_cpu_offload = True
16)
17
18import torch
19hf_auth = "YOUR-HUGGINGFACE-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/llama-2-medical-consultation")
35model = PeftModel.from_pretrained(model, "Ashishkr/llama-2-medical-consultation").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
45def llama_generate(
46 model: AutoModelForCausalLM,
47 tokenizer: AutoTokenizer,
48 prompt: str,
49 max_new_tokens: int = 128,
50 temperature: float = 0.92):
51
52 inputs = tokenizer(
53 [prompt],
54 return_tensors="pt",
55 return_token_type_ids=False,
56 ).to(
57 device
58 )
59
60 # Check if bfloat16 is supported, otherwise use float16
61 dtype_to_use = torch.float32
62 with torch.autocast("cuda", dtype=dtype_to_use):
63 response = model.generate(
64 **inputs,
65 max_new_tokens=max_new_tokens,
66 temperature=temperature,
67 return_dict_in_generate=True,
68 eos_token_id=tokenizer.eos_token_id,
69 pad_token_id=tokenizer.pad_token_id,
70 )
71
72 decoded_output = tokenizer.decode(
73 response["sequences"][0],
74 skip_special_tokens=True,
75 )
76
77 return decoded_output[len(prompt) :]
78
79prompt = """
80 instruction: "If you are a doctor, please answer the medical questions based on the patient's description." \n
81
82input: "Hi, I had a subarachnoid bleed and coiling of brain aneurysm last year.
83I am having some major bilateral temple pain along with numbness that comes and
84goes in my left arm/hand/fingers. I have had headaches since the aneurysm,
85but this is different. Also, my moods have been horrible for the past few weeks.\n
86
87response: """
88# You can use the function as before
89response = llama_generate(
90 model,
91 tokenizer,
92 prompt,
93 max_new_tokens=100,
94 temperature=0.92,
95)
96
97print(response)
98