Update README.md

README.md

@@ -21,12 +21,18 @@ Bulgarian, Chinese, Czech, Dutch, English, Estonian, Finnish, French, German, Gr
 Introducing a cutting-edge model tailored to the task of extracting entities from sensitive text and anonymizing it. This model specializes in identifying and safeguarding confidential information, ensuring organizations' compliance with stringent data privacy regulations and minimizing the potential for inadvertent disclosure of classified data and trade secrets.
 # Example Usage
 ```python
+!pip install sentencepiece
+!pip install transformers
+```
+```python
 import torch
+import json
+import re
 from transformers import AutoTokenizer, AutoModelForCausalLM
+
 tokenizer = AutoTokenizer.from_pretrained("metricspace/EntityAnonymization-3B-V0.9")
 model = AutoModelForCausalLM.from_pretrained("metricspace/EntityAnonymization-3B-V0.9", torch_dtype=torch.bfloat16)
-
-import re
+model.to("cuda:0")
 
 def extract_last_assistant_response(input_text):
     # Find the occurrence of "ASSISTANT:" in the input text
@@ -37,39 +43,132 @@ def extract_last_assistant_response(input_text):
     response = input_text[start_index:].strip()
     return response
 
+# Input example
+text_to_anonymize = '''Subject: HR Incident Report: Speculation of Drug Misuse by Mr. Benjamin Mitchell
 
-text_to_anonymize = '''Our organization manages a sophisticated data analytics platform ([login to view URL]) that highlights our cutting-edge data visualization techniques. In response to evolving business needs, we've recognized the imperative to optimize our data handling processes. As part of this initiative, we're seizing the opportunity to standardize the codebase for our web and mobile applications using a unified approach with Vue.js.
-We're currently seeking a talented developer to spearhead this transformation, ensuring a seamless separation between backend data processing and frontend presentation layers. The revised architecture will incorporate three critical APIs (Google Maps for location services, AccuWeather for weather data, and our in-house Analytica API for advanced analytics).
-The backend restructuring is a critical component, designed to serve as a showcase for the capabilities of our Analytica API. The frontend, both for the web and mobile interfaces, will maintain the current user experience using the existing design assets.
-We are actively searching for a Vue.js developer who can efficiently interpret our project vision and deliver an elegant, sustainable solution.'''
+Dear Mrs. Alice Williams,
 
+I trust you're well. I wish to bring to your attention a concerning matter involving one of our esteemed employees, Mr. Benjamin Mitchell.
 
-prompt = f'USER: Resample the entities: {text_to_anonymize}\n\nASSISTANT:'
-inputs = tokenizer(prompt, return_tensors='pt').to('cuda')
-output_entities = model.generate(inputs.input_ids, max_new_tokens=250, do_sample=False, top_k=50, top_p=0.98, num_beams=1)
-output_entities_text = tokenizer.decode(output_entities[0], skip_special_tokens=True)
+Employee Details:
 
-# extracting entities text from assistant response
-generated_part = extract_assistant_response(output_text_1)
+Name: Benjamin Mitchell
+Position: Senior Marketing Creative
+Department: Marketing
+Date of Joining: January 15, 2020
+Reporting Manager: Mrs. Jane Fitzgerald
 
-prompt_2 = f"USER: Rephrase with {generated_part}: {text_to_anonymize}\n\nASSISTANT:"
-inputs = tokenizer(prompt_2, return_tensors='pt').to('cuda')
-output_resampled = model.generate(inputs.input_ids, max_new_tokens=500, do_sample=False, top_k=50, top_p=0.98)
-output_resampled_text = tokenizer.decode(output_resampled[0], skip_special_tokens=True)
+Incident Details:
+Date: October 25, 2023
+Location: Restroom, 4th Floor
+Time: 11:45 AM
+
+Description of Incident:
+On 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.
+
+Witness Accounts:
+Ms. Emily Clark: "Benjamin seemed distracted and not his usual self today. He's been taking frequent breaks and appears a bit disoriented."
+Mr. Robert Taylor: "I found some pills near his chair on the floor. It's concerning, and I felt it necessary to report."
 
+Immediate Actions Taken:
+Mr. Benjamin Mitchell was approached by HR for a preliminary conversation to understand the situation.
+Mrs. Jane Fitzgerald, his reporting manager, was made aware of the concerns.
 
-print(output_resampled_text)
+Recommendations:
+It's crucial to have a private and supportive conversation with Mr. Mitchell to understand if there's an underlying issue.
+Consider referring Benjamin to our Employee Assistance Program (EAP) for counseling or support.
+It may be beneficial to organize a session on drug awareness and workplace safety for all employees.
+It'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.
 
+Looking forward to your direction on this matter.
+'''
+print(text_to_anonymize)
+
+# Step 1: Extracting entities from text
+prompt = f'USER: Resample the entities: {text_to_anonymize}\n\nASSISTANT:'
+inputs = tokenizer(prompt, return_tensors='pt').to('cuda:0')
+output_entities = model.generate(inputs.input_ids, max_new_tokens=250, do_sample=True, temperature=0.85)
+raw_output_entities_text = tokenizer.decode(output_entities[0])
+entities = extract_last_assistant_response(raw_output_entities_text)
+
+print('-----------Entities----------------')
+try:
+  entities = re.search(r"\{.*?\}", entities, re.DOTALL).group(0)
+  data_dict = eval(entities)
+  formatted_json = json.dumps(data_dict, indent=4)
+  print(formatted_json)
+except:
+  #bad formated json
+  print(entities)
 #output
 '''
-Our enterprise manages an advanced data analysis platform ([login to view URL]) that highlights our innovative data visualization methods. In response to evolving business needs, we've recognized the imperative to optimize our data handling processes. As part of this initiative, we're seizing the opportunity to standardize the codebase for our online and mobile applications using a unified approach with Vega.js.
-We're currently seeking a talented developer to spearhead this transformation, ensuring a seamless separation between backend data processing and frontend presentation layers. The revised architecture will incorporate three critical APIs (Maple Maps for location services, MeteorWeather for weather data, and our in-house Analytica API for advanced analytics).
-The backend restructuring is a critical component, designed to serve as a showcase for the capabilities of our Analytica API. The frontend, both for the web and mobile interfaces, will maintain the current user experience using the existing design assets.
-We are actively searching for a Vega.js developer who can efficiently interpret our project vision and deliver an elegant, sustainable solution
+{
+    "Mr. Benjamin Mitchell": "Mr. Edward Martin",
+    "Mrs. Alice Williams": "Mrs. Charlotte Johnson",
+    "January 15, 2020": "January 15, 2020",
+    "Mrs. Jane Fitzgerald": "Mrs. Jane Anderson",
+    "October 25, 2023": "October 25, 2023",
+    "4th Floor": "topmost floor",
+    "11:45 AM": "midday",
+    "Emily Clark": "Marie Foster",
+    "Employee Assistance Program (EAP)": "Personal Assistance Program (PAP)",
+    "Robert Taylor": "Benjamin Adams",
+}
 '''
+
+# Step 2: Use entities to resample the original text
+prompt_2 = f"USER: Rephrase with {entities}: {text_to_anonymize}\n\nASSISTANT:"
+inputs = tokenizer(prompt_2, return_tensors='pt').to('cuda:0')
+output_resampled = model.generate(inputs.input_ids, max_length=2048)
+raw_output_resampled_text = tokenizer.decode(output_resampled[0])
+resampled_text = extract_last_assistant_response(raw_output_resampled_text)
+print('---------Anonymized Version--------')
+print(resampled_text)
+#output:
+'''
+Subject: HR Incident Report: Speculation of Drug Misuse by Mr. Edward Martin
+
+Dear Mrs. Charlotte Johnson,
+
+I trust you're well. I wish to bring to your attention a concerning matter involving one of our esteemed employees, Mr. Edward Martin.
+
+Employee Details:
+
+Name: Edward Martin
+Position: Senior Marketing Creative
+Department: Marketing
+Date of Joining: January 15, 2020
+Reporting Manager: Mrs. Jane Anderson
+
+Incident Details:
+Date: October 25, 2023
+Location: Restroom, topmost floor
+Time: midday
+
+Description of Incident:
+On 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.
+
+Witness Accounts:
+Ms. Marie Foster: "Edward seemed distracted and not his usual self today. He's been taking frequent breaks and appears a bit disoriented."
+Mr. Benjamin Adams: "I found some pills near his chair on the floor. It's concerning, and I felt it necessary to report."
+
+Immediate Actions Taken:
+Mr. Edward Martin was approached by People Management for a preliminary conversation to understand the situation.
+Mrs. Jane Anderson, his reporting manager, was made aware of the concerns.
+
+Recommendations:
+It's crucial to have a private and supportive conversation with Mr. Martin to understand if there's an underlying issue.
+Consider referring Edward to our Personal Assistance Program (PAP) for counseling or support.
+It may be beneficial to organize a session on drug awareness and workplace safety for all employees.
+It'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.
+
+Looking forward to your direction on this matter.
+'''
+
 ```
 
 
+
 # Example inverted usage
 ```python
 import torch