luispoveda93/MiniCPM4-0.5B-PII-tagger-lora

MiniCPM4-0.5B PII Tagger (LoRA, full 100K)

A LoRA adapter on openbmb/MiniCPM4-0.5B that turns the model into a token-classification PII/PHI detector, trained on the full 100K-example train split of nvidia/Nemotron-PII, with the same recipe as the MiniCPM5-1B run (1B full).

Important: custom head + trust_remote_code

MiniCPM is not a native transformers architecture, and the remote code in the base repo has no AutoModelForTokenClassification entry. This adapter therefore lives on top of a small wrapper: the MiniCPMModel backbone (loaded via trust_remote_code=True) plus a linear score head (nn.Linear(1024, 55, bias=False), mirroring the convention of the repo's own MiniCPMForSequenceClassification). The head is saved inside the adapter via PEFT's modules_to_save=["score"]. Use the loading code below verbatim.

Results

Evaluated on 2,000 held-out examples from the nvidia/Nemotron-PII test split (seed 42 shuffle), same protocol as the MiniCPM5-1B runs (exact token/span match; spans are maximal runs of consecutive tokens with the same non-O label):

Metric	MiniCPM4-0.5B (this)	MiniCPM5-1B (full 100K)
Token P / R / F1	0.9722 / 0.9701 / 0.9711	0.9741 / 0.9704 / 0.9723
Span P / R / F1	0.8602 / 0.9091 / 0.8840	0.8942 / 0.9225 / 0.9081
eval_loss	0.0355	0.0301
Peak GPU memory	17.1 GB	22.4 GB
Train runtime (1×A10G)	~2 h 03 m	~2 h 41 m

The 0.5B model nearly matches the 1B on token-level F1 (−0.1 points) and gives up ~2.4 points of span F1 — with a third of the memory footprint and less training time.

Usage

import torch
from torch import nn
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.modeling_outputs import TokenClassifierOutput
from peft import PeftModel, LoraConfig, get_peft_model, TaskType

REPO = "luispoveda93/MiniCPM4-0.5B-PII-tagger-lora"
BASE = "openbmb/MiniCPM4-0.5B"

class MiniCPMTokenClassifier(nn.Module):
    def __init__(self, backbone, config, num_labels):
        super().__init__()
        self.config = config
        self.num_labels = num_labels
        self.backbone = backbone
        self.dropout = nn.Dropout(0.1)
        self.score = nn.Linear(config.hidden_size, num_labels, bias=False)
        self.score.to(next(backbone.parameters()).dtype)
        nn.init.normal_(self.score.weight, mean=0.0, std=0.02)
    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
        out = self.backbone(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
        h = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
        logits = self.score(self.dropout(h))
        loss = None
        if labels is not None:
            loss = nn.functional.cross_entropy(logits.view(-1, self.score.out_features).float(), labels.view(-1), ignore_index=-100)
        return TokenClassifierOutput(loss=loss, logits=logits)

causal = AutoModelForCausalLM.from_pretrained(BASE, trust_remote_code=True, dtype=torch.bfloat16)
causal.config.num_labels = 55
wrapper = MiniCPMTokenClassifier(causal.model, causal.config, 55)
lora = LoraConfig(task_type=TaskType.TOKEN_CLS, r=16, lora_alpha=32, lora_dropout=0.1,
                  target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
                  modules_to_save=["score"])
model = PeftModel.from_pretrained(wrapper, REPO)
tok = AutoTokenizer.from_pretrained(REPO)
if tok.pad_token_id is None:
    tok.pad_token = tok.eos_token

labels = ["account_number", "age", "api_key", "bank_routing_number", "biometric_identifier",
"blood_type", "certificate_license_number", "city", "company_name", "coordinate", "country",
"county", "credit_debit_card", "customer_id", "cvv", "date", "date_of_birth", "date_time",
"device_identifier", "education_level", "email", "employee_id", "employment_status", "fax_number",
"first_name", "gender", "health_plan_beneficiary_number", "http_cookie", "ipv4", "ipv6",
"language", "last_name", "license_plate", "mac_address", "medical_record_number", "national_id",
"occupation", "password", "phone_number", "pin", "political_view", "postcode", "race_ethnicity",
"religious_belief", "sexuality", "ssn", "state", "street_address", "swift_bic", "tax_id", "time",
"unique_id", "url", "user_name", "vehicle_identifier"]
id2label = {i: l for i, l in enumerate(labels)}

text = "My date of birth is 1987-05-22 and I live at 87 Avenida De La Estrella."
enc = tok(text, return_tensors="pt", return_offsets_mapping=True, truncation=True, max_length=1024)
offsets = enc.pop("offset_mapping")[0].tolist()
with torch.no_grad():
    preds = model(**enc).logits.argmax(-1)[0].tolist()
for (s, e), p in zip(offsets, preds):
    if p != 0 and not (s == 0 and e == 0):
        print(text[s:e], "->", id2label[p])

Adjacent predicted tokens with the same label form one entity span. Requires transformers>=4.53 (the base repo's remote code imports CacheLayerMixin) and peft.

Training details

• Base: openbmb/MiniCPM4-0.5B (MiniCPMForCausalLM remote code, MiniCPMModel backbone + score head, 55 labels, bf16)
• Data: all 100,000 examples of nvidia/Nemotron-PII default/train after a seed-42 shuffle; spans mapped to tokens via fast-tokenizer offset_mapping (IO scheme, -100 on special tokens); max_length=1024
• LoRA: r=16, alpha=32, dropout=0.1, target modules q/k/v/o/gate/up/down_proj + modules_to_save=["score"] — 8.4M trainable params (1.9%)
• Hyperparameters: lr 2e-4 (cosine, no warmup), effective batch 32 (8 × grad-accum 4), 1 epoch (3,125 optimizer steps), bf16, seed 42
• Hardware: 1× A10G (24 GB), peak GPU memory 17.1 GB, train runtime ~2 h 03 min

Data caveats and limitations

• ~1.5% of spans in nvidia/Nemotron-PII have a char-level mismatch between spans[i].text and text[start:end]; char offsets were treated as ground truth for labeling.
• Single epoch, aggregate metrics only — rare categories likely underperform.
• Synthetic, persona-grounded text only; real-world documents may shift.
• Token-level boundaries are exact-match; adjacent same-label spans merge into one prediction.

Job provenance

Trained with HF Jobs (minicpm4-pii-full-100k-a10g, job 6aa1122d900620b5c77e5c36), A10G-small, 2026-09-09. MiniCPM5-1B comparison run: 6aa0708832d5d0c22c5ae6f2 (100K).