1import torch
2from peft import PeftModel
3from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
4
5model_name = "meta-llama/Llama-2-70b-chat-hf"
6adapters_name = 'wj2003/Pongo-70B'
7
8model = AutoModelForCausalLM.from_pretrained(
9 model_name,
10 load_in_4bit=True,
11 torch_dtype=torch.bfloat16,
12 device_map="auto",
13 max_memory={i: '48000MB' for i in range(torch.cuda.device_count())},
14 quantization_config=BitsAndBytesConfig(
15 load_in_4bit=True,
16 bnb_4bit_compute_dtype=torch.bfloat16,
17 bnb_4bit_use_double_quant=True,
18 bnb_4bit_quant_type='nf4'
19 ),
20)
21model = PeftModel.from_pretrained(model, adapters_name)
22tokenizer = AutoTokenizer.from_pretrained(adapters_name)
23prompt = "find potential security issues in the following code. If it has vulnerability, " \
24 "output: Vulnerabilities " \
25 "Detected: type of vulnerability. otherwise output<no vulnerability detected>.Here is the complete code: "
26
27# Provide your code
28code=""
29formatted_prompt = (
30 f"{prompt + code}"
31 )
32inputs = tokenizer(formatted_prompt,return_tensors="pt").to("cuda:0")
33outputs = model.generate(inputs=inputs.input_ids, max_new_tokens=1024)