PIIMask-NOR

PIIMask-NOR Model

The PIIMask-NOR model is a specialized language model fine-tuned for the task of Personal Identifiable Information (PII) redaction in Norwegian, Bokmål. It is based on the "google/gemma-1.1-2b-it" model and trained to identify and redact various types of PII in text while maintaining the grammatical structure of sentences.

Model Description

• Model Name: PIIMask-NOR
• Base Model: google/gemma-1.1-2b-it
• Quantization: 4-bit quantization using NF4 with double quantization and float16 compute dtype.
• Training Steps: The model checkpoints are available at 1, 2, 3, and 4 epochs.

Methodology

The PIIMask-NOR model was fine-tuned using the ai4privacy/pii-masking-65k dataset, which was machine translated into Norwegian, Bokmål. The training process involved several epochs to improve the model's ability to accurately redact PII from text. The quantization configuration was applied to make the model more efficient for deployment.

Usage

Installation

To use the PIIMask-NOR model, you need to have the transformers and datasets libraries installed. You can install them using pip:

pip install transformers datasets

Code Example

Here is a code example to load and use the PIIMask-NOR model for PII redaction:

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
import torch

# Quantization configuration
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
)

# System instructions for PII redaction
system_instructions = """Erstatt følgende typer personopplysninger i teksten nedenfor med '[REDACTED]': [FIRST_NAME_x], [CITY_x], [COUNTRY_x]. Sørg for at hver type informasjon erstattes på en måte som opprettholder den grammatiske strukturen i setningen. Du skal kun returnere den nye teksten med de relevante erstatningene utført, uten den opprinnelige teksten eller noen tilleggsannotasjoner.
Input:"""

example_prompt = "Jeg heter Clara og bor i Bergen, Norge."

# Load model function
def load_model(repo, step):
    model = AutoModelForCausalLM.from_pretrained(repo,
                                                 device_map="cuda:0",
                                                 trust_remote_code=True,
                                                 quantization_config=bnb_config,
                                                 adapter_kwargs={"subfolder": f"checkpoint-{step}"},
                                                 attn_implementation="flash_attention_2")
    return model

# Initialize tokenizer and model
device = "cuda"
tokenizer = AutoTokenizer.from_pretrained("google/gemma-1.1-2b-it", use_fast=True)

# Apply chat template for input
chat = [
    {"role": "system", "content": system_instructions},
    {"role": "user", "content": example_prompt},
]

inputs = tokenizer.apply_chat_template(chat, tokenize=False, return_tensors="pt", padding=True, truncation=False)
model = load_model("UlrikKoren/PIIMask-NOR", step=1159)
outputs = model.generate(input_ids=inputs['input_ids'].to(device), max_new_tokens=2048)
decoded_outputs = [tokenizer.decode(output, skip_special_tokens=False) for output in outputs]
print(decoded_outputs[0])

Checkpoints

The model checkpoints for different training epochs can be accessed as follows:

• Epoch 1: UlrikKoren/PIIMask-NOR/tree/main/checkpoint-579
• Epoch 2: UlrikKoren/PIIMask-NOR/checkpoint-1159
• Epoch 3: UlrikKoren/PIIMask-NOR/checkpoint-1739
• Epoch 4: UlrikKoren/PIIMask-NOR/checkpoint-2316

Compliance with Gemma Terms of Use

This model is a derivative of the "google/gemma-1.1-2b-it" model and complies with the Gemma Terms of Use:

• Distribution: Any distribution of this model or its derivatives must include the use restrictions specified in the Gemma Terms of Use and provide notice to subsequent users.
• Notices: The model is distributed with the following notice: “Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms”.
• Modifications: Any modified files carry prominent notices stating the modifications made.
• Prohibited Uses: The use of this model is subject to the restrictions outlined in the Gemma Prohibited Use Policy.
• Trademarks: This distribution does not grant any rights to use Google’s trademarks, trade names, or logos.

License

The PIIMask-NOR model is distributed under the same terms as the base model. For more details, please refer to the Gemma Terms of Use.