Views
No views yet
1
2import transformers
3from peft import PeftModel, PeftConfig
4from transformers import AutoModelForCausalLM, AutoTokenizer
5import torch
6from torch import cuda, bfloat16
7
8base_model_id = 'meta-llama/Llama-2-7b-hf'
9
10device = f'cuda:{cuda.current_device()}' if cuda.is_available() else 'cpu'
11
12bnb_config = transformers.BitsAndBytesConfig(
13 load_in_4bit=True,
14 bnb_4bit_quant_type='nf4',
15 bnb_4bit_use_double_quant=True,
16 bnb_4bit_compute_dtype=bfloat16
17)
18
19
20hf_auth = "hf_your-huggingface-access-token"
21model_config = transformers.AutoConfig.from_pretrained(
22 base_model_id,
23 use_auth_token=hf_auth
24)
25
26model = transformers.AutoModelForCausalLM.from_pretrained(
27 base_model_id,
28 trust_remote_code=True,
29 config=model_config,
30 quantization_config=bnb_config,
31 device_map='auto',
32 use_auth_token=hf_auth
33)
34
35config = PeftConfig.from_pretrained("Ashishkr/PII-Masking")
36model = PeftModel.from_pretrained(model, "Ashishkr/PII-Masking").to(device)
37
38model.eval()
39print(f"Model loaded on {device}")
40
41tokenizer = transformers.AutoTokenizer.from_pretrained(
42 base_model_id,
43 use_auth_token=hf_auth
44)
451def remove_pii_info(
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).to(device)
12
13 max_new_tokens = inputs["input_ids"].shape[1]
14
15 # Check if bfloat16 is supported, otherwise use float16
16 dtype_to_use = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
17
18 with torch.autocast("cuda", dtype=dtype_to_use):
19 response = model.generate(
20 **inputs,
21 max_new_tokens=max_new_tokens,
22 temperature=temperature,
23 return_dict_in_generate=True,
24 eos_token_id=tokenizer.eos_token_id,
25 pad_token_id=tokenizer.pad_token_id,
26 )
27
28 decoded_output = tokenizer.decode(
29 response["sequences"][0],
30 skip_special_tokens=True,
31 )
32
33 return decoded_output[len(prompt) :]
34
35prompt = """
36 Input: "John Doe, currently lives at 1234 Elm Street, Springfield, Anywhere 12345.
37 He can be reached at johndoe@email.com or at the phone number 555-123-4567. His social security number is 123-45-6789,
38 and he has a bank account number 9876543210 at Springfield Bank. John attended Springfield University where he earned
39 a Bachelor's degree in Computer Science. He now works at Acme Corp and his employee ID is 123456. John's medical record number
40 is MRN-001234, and he has a history of asthma and high blood pressure. His primary care physician is Dr. Jane Smith,
41 who practices at Springfield Medical Center. His recent blood test results show a cholesterol level of 200 mg/dL and a
42 blood glucose level of 90 mg/dL.
43" Output: """
44# You can use the function as before
45response = remove_pii_info(
46 model,
47 tokenizer,
48 prompt,
49 temperature=0.7)
50
51print(response)