Views
No views yet
1!pip install sentencepiece
2!pip install transformers1import torch
2import json
3import re
4from transformers import AutoTokenizer, AutoModelForCausalLM
5
6tokenizer = AutoTokenizer.from_pretrained("metricspace/EntityAnonymization-3B-V0.9")
7model = AutoModelForCausalLM.from_pretrained("metricspace/EntityAnonymization-3B-V0.9", torch_dtype=torch.bfloat16)
8model.to("cuda:0")
9
10def extract_last_assistant_response(input_text):
11 # Find the occurrence of "ASSISTANT:" in the input text
12 match = re.search(r'ASSISTANT:', input_text)
13
14 # Get the index where the last "ASSISTANT:" ends
15 start_index = match.end()
16 response = input_text[start_index:].strip()
17 return response
18
19# Input example
20text_to_anonymize = '''Subject: HR Incident Report: Speculation of Drug Misuse by Mr. Benjamin Mitchell
21
22Dear Mrs. Alice Williams,
23
24I trust you're well. I wish to bring to your attention a concerning matter involving one of our esteemed employees, Mr. Benjamin Mitchell.
25
26Employee Details:
27
28Name: Benjamin Mitchell
29Position: Senior Marketing Creative
30Department: Marketing
31Date of Joining: January 15, 2020
32Reporting Manager: Mrs. Jane Fitzgerald
33
34Incident Details:
35Date: October 25, 2023
36Location: Restroom, 4th Floor
37Time: 11:45 AM
38
39Description of Incident:
40On the date specified, a few colleagues reported unusual behavior exhibited by Mr. Mitchell, which raised concerns about potential drug misuse. Witnesses mentioned that Benjamin appeared disoriented and was found in the restroom for an extended period. Some employees also discovered unidentified pills in close proximity to his chair.
41
42Witness Accounts:
43Ms. Emily Clark: "Benjamin seemed distracted and not his usual self today. He's been taking frequent breaks and appears a bit disoriented."
44Mr. Robert Taylor: "I found some pills near his chair on the floor. It's concerning, and I felt it necessary to report."
45
46Immediate Actions Taken:
47Mr. Benjamin Mitchell was approached by HR for a preliminary conversation to understand the situation.
48Mrs. Jane Fitzgerald, his reporting manager, was made aware of the concerns.
49
50Recommendations:
51It's crucial to have a private and supportive conversation with Mr. Mitchell to understand if there's an underlying issue.
52Consider referring Benjamin to our Employee Assistance Program (EAP) for counseling or support.
53It may be beneficial to organize a session on drug awareness and workplace safety for all employees.
54It's of utmost importance to handle this situation with sensitivity and discretion, ensuring the wellbeing of Mr. Mitchell and maintaining the integrity of our workplace environment. This email serves as a formal documentation of the incident. We'll determine the subsequent course of action based on your guidance and the recommendations provided.
55
56Looking forward to your direction on this matter.
57'''
58print(text_to_anonymize)
59
60# Step 1: Extracting entities from text
61prompt = f'USER: Resample the entities: {text_to_anonymize}\n\nASSISTANT:'
62inputs = tokenizer(prompt, return_tensors='pt').to('cuda:0')
63output_entities = model.generate(inputs.input_ids, max_new_tokens=300, do_sample=False, temperature=0.8, penalty_alpha=1.3, top_k=180, num_beams=5, repetition_penalty=2.3)
64
65raw_output_entities_text = tokenizer.decode(output_entities[0])
66entities = extract_last_assistant_response(raw_output_entities_text)
67
68print('-----------Entities----------------')
69try:
70 entities = re.search(r"\{.*?\}", entities, re.DOTALL).group(0)
71 data_dict = eval(entities)
72 formatted_json = json.dumps(data_dict, indent=4)
73 print(formatted_json)
74except:
75 #bad formated json
76 print(entities)
77#output
78'''
79{
80 "Mr. Benjamin Mitchell": "Mr. Edward Martin",
81 "Mrs. Alice Williams": "Mrs. Charlotte Johnson",
82 "January 15, 2020": "January 15, 2020",
83 "Mrs. Jane Fitzgerald": "Mrs. Jane Anderson",
84 "October 25, 2023": "October 25, 2023",
85 "4th Floor": "topmost floor",
86 "11:45 AM": "midday",
87 "Emily Clark": "Marie Foster",
88 "Employee Assistance Program (EAP)": "Personal Assistance Program (PAP)",
89 "Robert Taylor": "Benjamin Adams",
90}
91'''
92
93# Step 2: Use entities to resample the original text
94prompt_2 = f"USER: Rephrase with {entities}: {text_to_anonymize}\n\nASSISTANT:"
95inputs = tokenizer(prompt_2, return_tensors='pt').to('cuda:0')
96output_resampled = model.generate(inputs.input_ids, max_length=2048)
97raw_output_resampled_text = tokenizer.decode(output_resampled[0])
98resampled_text = extract_last_assistant_response(raw_output_resampled_text)
99print('---------Anonymized Version--------')
100print(resampled_text)
101#output:
102'''
103Subject: HR Incident Report: Speculation of Drug Misuse by Mr. Edward Martin
104
105Dear Mrs. Charlotte Johnson,
106
107I trust you're well. I wish to bring to your attention a concerning matter involving one of our esteemed employees, Mr. Edward Martin.
108
109Employee Details:
110
111Name: Edward Martin
112Position: Senior Marketing Creative
113Department: Marketing
114Date of Joining: January 15, 2020
115Reporting Manager: Mrs. Jane Anderson
116
117Incident Details:
118Date: October 25, 2023
119Location: Restroom, topmost floor
120Time: midday
121
122Description of Incident:
123On the date specified, a few colleagues reported unusual behavior exhibited by Mr. Martin, which raised concerns about potential drug misuse. Witnesses mentioned that Edward appeared disoriented and was found in the restroom for an extended period. Some employees also discovered unidentified pills in close proximity to his chair.
124
125Witness Accounts:
126Ms. Marie Foster: "Edward seemed distracted and not his usual self today. He's been taking frequent breaks and appears a bit disoriented."
127Mr. Benjamin Adams: "I found some pills near his chair on the floor. It's concerning, and I felt it necessary to report."
128
129Immediate Actions Taken:
130Mr. Edward Martin was approached by People Management for a preliminary conversation to understand the situation.
131Mrs. Jane Anderson, his reporting manager, was made aware of the concerns.
132
133Recommendations:
134It's crucial to have a private and supportive conversation with Mr. Martin to understand if there's an underlying issue.
135Consider referring Edward to our Personal Assistance Program (PAP) for counseling or support.
136It may be beneficial to organize a session on drug awareness and workplace safety for all employees.
137It's of utmost importance to handle this situation with sensitivity and discretion, ensuring the wellbeing of Mr. Martin and maintaining the integrity of our workplace environment. This email serves as a formal documentation of the incident. We'll determine the subsequent course of action based on your guidance and the recommendations provided.
138
139Looking forward to your direction on this matter.
140'''
1411import torch
2import json
3import re
4from transformers import AutoTokenizer, AutoModelForCausalLM
5
6tokenizer = AutoTokenizer.from_pretrained("metricspace/EntityAnonymization-3B-V0.9")
7model = AutoModelForCausalLM.from_pretrained("metricspace/EntityAnonymization-3B-V0.9", torch_dtype=torch.bfloat16)
8model.to("cuda:0")
9
10
11# Anonymized input
12anonymized_text = '''Subject: HR Incident Report: Speculation of Drug Misuse by Mr. Edward Martin
13
14Dear Mrs. Charlotte Johnson,
15
16I trust you're well. I wish to bring to your attention a concerning matter involving one of our esteemed employees, Mr. Edward Martin.
17
18Employee Details:
19
20Name: Edward Martin
21Position: Senior Marketing Creative
22Department: Marketing
23Date of Joining: January 15, 2020
24Reporting Manager: Mrs. Jane Anderson
25
26Incident Details:
27Date: October 25, 2023
28Location: Restroom, topmost floor
29Time: midday
30
31Description of Incident:
32On the date specified, a few colleagues reported unusual behavior exhibited by Mr. Martin, which raised concerns about potential drug misuse. Witnesses mentioned that Edward appeared disoriented and was found in the restroom for an extended period. Some employees also discovered unidentified pills in close proximity to his chair.
33
34Witness Accounts:
35Ms. Marie Foster: "Edward seemed distracted and not his usual self today. He's been taking frequent breaks and appears a bit disoriented."
36Mr. Benjamin Adams: "I found some pills near his chair on the floor. It's concerning, and I felt it necessary to report."
37
38Immediate Actions Taken:
39Mr. Edward Martin was approached by People Management for a preliminary conversation to understand the situation.
40Mrs. Jane Anderson, his reporting manager, was made aware of the concerns.
41
42Recommendations:
43It's crucial to have a private and supportive conversation with Mr. Martin to understand if there's an underlying issue.
44Consider referring Edward to our Personal Assistance Program (PAP) for counseling or support.
45It may be beneficial to organize a session on drug awareness and workplace safety for all employees.
46It's of utmost importance to handle this situation with sensitivity and discretion, ensuring the wellbeing of Mr. Martin and maintaining the integrity of our workplace environment. This email serves as a formal documentation of the incident. We'll determine the subsequent course of action based on your guidance and the recommendations provided.
47
48Looking forward to your direction on this matter.
49'''
50
51
52# Entities map
53
54entities_map = '''
55{
56 "Mr. Benjamin Mitchell": "Mr. Edward Martin",
57 "Mrs. Alice Williams": "Mrs. Charlotte Johnson",
58 "January 15, 2020": "January 15, 2020",
59 "Mrs. Jane Fitzgerald": "Mrs. Jane Anderson",
60 "October 25, 2023": "October 25, 2023",
61 "4th Floor": "topmost floor",
62 "11:45 AM": "midday",
63 "Emily Clark": "Marie Foster",
64 "Employee Assistance Program (EAP)": "Personal Assistance Program (PAP)",
65 "Robert Taylor": "Benjamin Adams",
66}
67'''
68
69
70
71# Step 1: Processing anonymized text with GPT-4
72
73import openai
74openai.api_key = f"<API_KEY>"
75
76completion = openai.ChatCompletion.create(
77 model="gpt-4",
78 messages=[
79 {"role": "user", "content": f"Write an official warning letter to the employee, that we do not tolerate it and with the next incident he gets fired in the name of Mrs. Charlotte Johnson Human Resources Manager. Here is the report with information. {anonymized_info}"}
80 ]
81)
82
83print(completion.choices[0].message.content)
84#output
85'''
86Subject: Official Warning – Substance Misuse Policy Violation
87
88Dear Mr. Edward Martin,
89
90We hope this letter finds you well. This letter serves as an official warning in regards to an incident that occurred on October 25, 2023, in which signs of drug misuse were reported. This alleged conduct is highly alarming and contrary to our company’s policies and guidelines.
91
92The incident detailed allegations of unusual behavior indicative of possible substance abuse. Colleagues reported that you appeared disoriented and were found in the restroom for an extended period. Additionally, unidentified pills were discovered near your workspace.
93
94Our company is committed to providing a safe and non-detrimental work environment for all its workforce. This commitment is compromised when any type of drug misuse occurs. We draw your attention to our Employee Handbook, specifically 'Section 5: Substance Misuse', that states any illegal drug use, substance misuse or distribution thereof is strictly prohibited and could be subject to severe disciplinary action, including termination.
95
96This is an official warning that such behavior misaligns with our workplace norms and cannot be tolerated. Another incident like this or similar breach of company guidelines may lead to further disciplinary action, up to and including termination of employment.
97
98Please note that this is not an assumption of your guilt but an assertion to remain vigilant against any practices that could harm you or the workplace environment. We encourage you to utilize our Personal Assistance Program (PAP) as a tool for counseling and support, if needed.
99
100We believe in your potential to rectify this situation and to maintain the high standards we are all accustomed to in our organization.
101
102Should you need assistance or if you wish to discuss this matter further, please feel free to reach out to me. We appreciate your immediate attention to this important issue.
103
104Yours sincerely,
105
106Mrs. Charlotte Johnson
107Human Resources Manager
108'''
109
110
111# Step 2: Replace the entities back in processed by GPT-4 text.
112
113import ast
114
115def swap_keys_and_values_in_string(input_str):
116 # Convert the input string to a dictionary
117 input_dict = ast.literal_eval(input_str)
118
119 # Swap the keys and values
120 swapped_dict = {v: k for k, v in input_dict.items()}
121
122 # Convert the swapped dictionary back to a string
123 swapped_str = str(swapped_dict)
124
125 return swapped_str
126
127
128gpt_response = completion.choices[0].message.content
129entities_map = swap_keys_and_values_in_string(entities_map)
130prompt = f"USER: Rephrase with {entities_map}: {gpt_response}\n\nASSISTANT:"
131inputs = tokenizer(prompt, return_tensors='pt').to('cuda:0')
132outputs = model.generate(inputs.input_ids, max_new_tokens=2048)
133output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
134
135#output:
136'''
137Subject: Official Warning – Substance Misuse Policy Violation
138
139Dear Mr. Benjamin Mitchell,
140
141We hope this letter finds you well. This letter serves as an official warning in regards to an incident that occurred on January 15,
142
1432020, in which signs of drug misuse were reported. This alleged conduct is highly alarming and contrary to our company’s policies and guidelines.
144
145The incident detailed allegations of unusual behavior indicative of possible substance abuse. Colleagues reported that you appeared disoriented and
146
147were found in the restroom for an extended period. Additionally, unidentified pills were discovered near your workspace.
148
149Our company is committed to providing a safe and non-detrimental work environment for all its workforce. This commitment is compromised when any
150
151type of drug misuse occurs. We draw your attention to our Employee Handbook, specifically 'Section 5: Substance Misuse', that states any illegal
152
153drug use, substance misuse or distribution thereof is strictly prohibited and could be subject to severe disciplinary action, including termination.
154
155This is an official warning that such behavior misaligns with our workplace norms and cannot be tolerated. Another incident like this or similar breach
156
157of company guidelines may lead to further disciplinary action, up to and including termination of employment.
158
159Please note that this is not an assumption of your guilt but an assertion to remain vigilant against any practices that could harm you or the workplace
160
161environment. We encourage you to utilize our Employee Assistance Program (EAP) as a tool for counseling and support, if needed.
162
163We believe in your potential to rectify this situation and to maintain the high standards we are all accustomed to in our organization.
164
165Should you need assistance or if you wish to discuss this matter further, please feel free to reach out to me. We appreciate your immediate attention
166
167to this important issue.
168
169Yours sincerely,
170
171Mrs. Alice Williams,
172
173Human Resources Manager.
174'''
175