deryaerman/mmbert_multilingual_pii_ner

import torch
import json
from transformers import AutoModel, AutoTokenizer
from torchcrf import CRF
import torch.nn as nn

class ModernBertCRF(nn.Module):
    def __init__(self, base_model_name, num_labels, id2label, label2id):
        super().__init__()
        self.num_labels = num_labels
        self.id2label = id2label
        self.label2id = label2id
        self.transformer = AutoModel.from_pretrained(base_model_name)
        hidden_size = self.transformer.config.hidden_size
        self.classifier = nn.Linear(hidden_size, num_labels)
        self.dropout = nn.Dropout(0.1)
        self.crf = CRF(num_labels, batch_first=True)

    def forward(self, input_ids, attention_mask, labels=None, **kwargs):
        kwargs.pop("token_type_ids", None)
        outputs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)
        sequence_output = self.dropout(outputs.last_hidden_state)
        emissions = self.classifier(sequence_output)
        if labels is not None:
            mask = attention_mask.bool()
            labels_for_crf = labels.clone()
            labels_for_crf[labels_for_crf == -100] = 0
            loss = -self.crf(emissions, labels_for_crf, mask=mask, reduction='mean')
            return {"loss": loss, "logits": emissions}
        else:
            return {"logits": emissions}

    def decode(self, emissions, mask):
        return self.crf.decode(emissions, mask=mask)

# Load model
model_dir = "deryaerman/mmbert_multilingual_pii_ner"

with open(f"{model_dir}/crf_config.json") as f:
    config = json.load(f)

model = ModernBertCRF(
    base_model_name=config["base_model_name"],
    num_labels=config["num_labels"],
    id2label=config["id2label"],
    label2id=config["label2id"],
)
model.load_state_dict(torch.load(f"{model_dir}/pytorch_model.bin", map_location="cpu"))
model.eval()

tokenizer = AutoTokenizer.from_pretrained(model_dir)
id2label = {int(k): v for k, v in config["id2label"].items()}

import re
import spacy

nlp = spacy.blank("en")          # use "de" for German, "xx" for multilingual
nlp.add_pipe("sentencizer")

def split_dialogue(text):
    """
    Split raw dialogue text into a list of sentences (each a list of tokens).
    Expects lines like: 'SPEAKER_00: Hello, my name is Peter.'
    """
    sentences = []
    for line in text.strip().splitlines():
        m = re.match(r"^(SPEAKER_\d+)\s*:\s*(.*)", line.strip())
        if m:
            line = m.group(2)
        if not line:
            continue
        doc = nlp(line)
        for sent in doc.sents:
            tokens = [tok.text for tok in sent if not tok.is_space]
            if tokens:
                sentences.append(tokens)
    return sentences

# Example
raw = """SPEAKER_00: Hello, my name is Peter.
SPEAKER_01: Hello, my name is Peter as well. Okay, and where do you come from? I come from Chicago."""

dialogue = split_dialogue(raw)
# [['Hello', ',', 'my', 'name', 'is', 'Peter', '.'],
#  ['Hello', ',', 'my', 'name', 'is', 'Peter', 'as', 'well', '.'],
#  ['Okay', ',', 'and', 'where', 'do', 'you', 'come', 'from', '?'],
#  ['I', 'come', 'from', 'Chicago', '.']]

def predict_sentences(sentences, model, tokenizer, id2label, device="cpu"):
    """
    sentences: list of sentences, each a list of word tokens
    Returns:   list of label lists, one per sentence
    """
    all_labels = []
    for tokens in sentences:
        enc = tokenizer(tokens, is_split_into_words=True,
                        return_tensors="pt", truncation=True, max_length=512).to(device)
        word_ids = enc.word_ids(batch_index=0)

        with torch.no_grad():
            outputs = model(**enc)
            emissions = outputs["logits"]
            mask = enc["attention_mask"].bool()
            preds = model.decode(emissions, mask)[0]

        word_labels = ["O"] * len(tokens)
        seen = set()
        for idx, wid in enumerate(word_ids):
            if wid is None or wid in seen:
                continue
            seen.add(wid)
            word_labels[wid] = id2label[preds[idx]]

        all_labels.append(word_labels)

    return all_labels


# Example: dialogue from above
results = predict_sentences(dialogue, model, tokenizer, id2label)

for sent_tokens, sent_labels in zip(dialogue, results):
    for token, label in zip(sent_tokens, sent_labels):
        if label != "O":
            print(f"{token:20s} -> {label}")

tokens = ["My", "name", "is", "John", "Smith", "and", "I", "live", "in", "Berlin", "."]

enc = tokenizer(tokens, is_split_into_words=True, return_tensors="pt", truncation=True, max_length=512)
word_ids = enc.word_ids(batch_index=0)

with torch.no_grad():
    outputs = model(**enc)
    emissions = outputs["logits"]
    mask = enc["attention_mask"].bool()
    preds = model.decode(emissions, mask)[0]

seen = set()
for idx, wid in enumerate(word_ids):
    if wid is None or wid in seen:
        continue
    seen.add(wid)
    label = id2label[preds[idx]]
    if label != "O":
        print(f"{tokens[wid]:20s} -> {label}")

@mastersthesis{erman2026multilingual,
  title={Multilingual De-Identification of Dialogue Data using Transformer-based NER},
  author={Erman, Derya},
  year={2026}
}