Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3import json
4
5# Load model and tokenizer
6model_name = "eternisai/Anonymizer-0.6B"
7tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
8model = AutoModelForCausalLM.from_pretrained(
9 model_name,
10 torch_dtype=torch.float16,
11 device_map="auto",
12 trust_remote_code=True
13)
14
15# Define the task instruction
16TASK_INSTRUCTION = """You are an anonymizer. Your task is to identify and replace personally identifiable information (PII) in the given text.
17Replace PII entities with semantically equivalent alternatives that preserve the context needed for a good response.
18If no PII is found or replacement is not needed, return an empty replacements list.
19
20REPLACEMENT RULES:
21• Personal names: Replace private or small-group individuals. Pick same culture + gender + era; keep surnames aligned across family members. DO NOT replace globally recognised public figures (heads of state, Nobel laureates, A-list entertainers, Fortune-500 CEOs, etc.).
22• Companies / organisations: Replace private, niche, employer & partner orgs. Invent a fictitious org in the same industry & size tier; keep legal suffix. Keep major public companies (anonymity set ≥ 1,000,000).
23• Projects / codenames / internal tools: Always replace with a neutral two-word alias of similar length.
24• Locations: Replace street addresses, buildings, villages & towns < 100k pop with a same-level synthetic location inside the same state/country. Keep big cities (≥ 1M), states, provinces, countries, iconic landmarks.
25• Dates & times: Replace birthdays, meeting invites, exact timestamps. Shift day/month by small amounts while KEEPING THE SAME YEAR to maintain temporal context. DO NOT shift public holidays or famous historic dates ("July 4 1776", "Christmas Day", "9/11/2001", etc.). Keep years, fiscal quarters, decade references unchanged.
26• Identifiers: (emails, phone #s, IDs, URLs, account #s) Always replace with format-valid dummies; keep domain class (.com big-tech, .edu, .gov).
27• Monetary values: Replace personal income, invoices, bids by × [0.8 – 1.25] to keep order-of-magnitude. Keep public list prices & market caps.
28• Quotes / text snippets: If the quote contains PII, swap only the embedded tokens; keep the rest verbatim."""
29
30# Define tool schema (required!)
31tools = [{
32 "type": "function",
33 "function": {
34 "name": "replace_entities",
35 "description": "Replace PII entities with anonymized versions",
36 "parameters": {
37 "type": "object",
38 "properties": {
39 "replacements": {
40 "type": "array",
41 "items": {
42 "type": "object",
43 "properties": {
44 "original": {"type": "string"},
45 "replacement": {"type": "string"}
46 },
47 "required": ["original", "replacement"]
48 }
49 }
50 },
51 "required": ["replacements"]
52 }
53 }
54}]
55
56# Your query to anonymize
57query = "Hi, my son Elijah works at TechStartup Inc and makes $85,000 per year."
58
59# Format messages properly (critical step!)
60messages = [
61 {"role": "system", "content": TASK_INSTRUCTION},
62 {"role": "user", "content": query + "\n/no_think"}
63]
64
65# Apply chat template with tools
66formatted_prompt = tokenizer.apply_chat_template(
67 messages,
68 tools=tools,
69 tokenize=False,
70 add_generation_prompt=True
71)
72
73# Tokenize and generate
74inputs = tokenizer(formatted_prompt, return_tensors="pt", truncation=True).to(model.device)
75outputs = model.generate(**inputs, max_new_tokens=250, temperature=0.3, do_sample=True, top_p=0.9)
76
77# Decode and extract response
78response = tokenizer.decode(outputs[0], skip_special_tokens=False)
79assistant_response = response.split("assistant")[-1].split("<|im_end|>")[0].strip()
80
81print("Response:", assistant_response)
82# Expected output format:
83# <|tool_call|>{"name": "replace_entities", "arguments": {"replacements": [{"original": "Elijah", "replacement": "Nathan"}, {"original": "TechStartup Inc", "replacement": "DataSoft LLC"}, {"original": "$85,000", "replacement": "$72,000"}]}}</|tool_call|>1def parse_replacements(response):
2 """Extract replacements from model response"""
3 try:
4 if '<|tool_call|>' in response:
5 start = response.find('<|tool_call|>') + len('<|tool_call|>')
6 end = response.find('</|tool_call|>')
7 elif '<tool_call>' in response:
8 start = response.find('<tool_call>') + len('<tool_call>')
9 end = response.find('</tool_call>')
10 else:
11 return None
12
13 if end != -1:
14 json_str = response[start:end].strip()
15 tool_data = json.loads(json_str)
16 return tool_data.get('arguments', {}).get('replacements', [])
17 except:
18 return None
19
20# Parse the response
21replacements = parse_replacements(assistant_response)
22if replacements:
23 for r in replacements:
24 print(f"Replace '{r['original']}' with '{r['replacement']}'")1<|tool_call|>
2{"name": "replace_entities", "arguments": {"replacements": [
3 {"original": "John", "replacement": "Marcus"},
4 {"original": "Microsoft", "replacement": "TechCorp"},
5 {"original": "$5000", "replacement": "$4200"}
6]}}
7</|tool_call|>1<|tool_call|>
2{"name": "replace_entities", "arguments": {"replacements": []}}
3</|tool_call|>tokenizer.apply_chat_template() with the tools parameter./no_think marker appended.<|tool_call|> tags (or <tool_call> in some versions).apply_chat_template with the tools parameter/no_think to the user query[BEGIN OF TASK INSTRUCTION]
You are an anonymizer. Your task is to identify and replace personally identifiable information (PII)...
[END OF TASK INSTRUCTION]
[BEGIN OF AVAILABLE TOOLS]
[{"type": "function", "function": {"name": "replace_entities", ...}}]
[END OF AVAILABLE TOOLS]
[BEGIN OF FORMAT INSTRUCTION]
Use the replace_entities tool to specify replacements...
[END OF FORMAT INSTRUCTION]
[BEGIN OF QUERY]
Your text to anonymize goes here
/no_think
[END OF QUERY]tokenizer.apply_chat_template() - never construct it manually.