Upload folder using huggingface_hub

2bf9c60 verified 5 months ago

3.31 kB

#!/usr/bin/env python3
"""Minimal runnable example for FrameByFrame/korean-pii-e5-base.
    pip install "transformers>=4.40" torch safetensors
    python usage.py
"""
import os, re
import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification

MODEL_ID = os.environ.get("MODEL_ID", "FrameByFrame/korean-pii-e5-base")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForTokenClassification.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
model.eval()
if torch.cuda.is_available():
    model.cuda()

_TRAILING_JOSA = ["이에요","이라고","입니다","이야","이랑","한테","에게","으로","이가","이는",
                  "에서","이고","예요","씨","님","이","가","은","는","을","를","야","아","에","의","랑","께","고"]
_DATE_END = re.compile(r".*(?:일|[0-9])", re.S)


def _normalize(text, label, s, e):
    while s < e and text[s] in " .,\t\n": s += 1
    while e > s and text[e - 1] in " .,\t\n": e -= 1
    if label == "private_date":
        m = _DATE_END.match(text[s:e])
        if m and m.end() > 0:
            e = s + m.end()
    elif label in ("private_person", "personal_handle", "private_address"):
        for _ in range(2):
            seg = text[s:e]
            for j in _TRAILING_JOSA:
                if seg.endswith(j) and (e - s) - len(j) >= 2:
                    e -= len(j); break
            else:
                break
    return s, e


def extract_pii(text, max_length=256):
    enc = tokenizer(text, truncation=True, max_length=max_length,
                    return_offsets_mapping=True, return_tensors="pt")
    offsets = enc.pop("offset_mapping")[0].tolist()
    with torch.no_grad():
        logits = model(**{k: v.to(model.device) for k, v in enc.items()}).logits
    pred = logits.argmax(-1)[0].tolist()
    id2label = model.config.id2label
    spans, active = [], None
    for i, lid in enumerate(pred):
        label = id2label[int(lid)]
        cs, ce = offsets[i]
        if cs == ce:
            if active: spans.append(active); active = None
            continue
        if label == "O":
            if active: spans.append(active); active = None
            continue
        prefix, cat = label.split("-", 1)
        if prefix in ("B", "S") or not active or active[0] != cat:
            if active: spans.append(active)
            active = [cat, cs, ce]
        else:
            active[2] = ce
    if active: spans.append(active)
    out = []
    for cat, s, e in spans:
        s, e = _normalize(text, cat, s, e)
        if text[s:e].strip():
            out.append({"label": cat, "start": s, "end": e, "text": text[s:e]})
    return out


def redact(text):
    spans = sorted(extract_pii(text), key=lambda s: s["start"], reverse=True)
    for s in spans:
        text = text[:s["start"]] + f"[{s['label'].upper()}]" + text[s["end"]:]
    return text


if __name__ == "__main__":
    for t in ["김민수님의 번호는 010-1234-5678입니다.",
              "계좌 110-234-567890으로 입금하고 minsu@example.com으로 알려주세요.",
              "이수진 고객님 생년월일은 1985년 3월 12일입니다."]:
        print(t)
        for sp in extract_pii(t):
            print(f"   {sp['label']:16} {sp['text']!r}")
        print("   REDACT:", redact(t)); print()

	#!/usr/bin/env python3
	"""Minimal runnable example for FrameByFrame/korean-pii-e5-base.

	pip install "transformers>=4.40" torch safetensors
	python usage.py
	"""
	import os, re
	import torch
	from transformers import AutoTokenizer, AutoModelForTokenClassification

	MODEL_ID = os.environ.get("MODEL_ID", "FrameByFrame/korean-pii-e5-base")
	tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
	model = AutoModelForTokenClassification.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
	model.eval()
	if torch.cuda.is_available():
	model.cuda()

	_TRAILING_JOSA = ["이에요","이라고","입니다","이야","이랑","한테","에게","으로","이가","이는",
	"에서","이고","예요","씨","님","이","가","은","는","을","를","야","아","에","의","랑","께","고"]
	_DATE_END = re.compile(r".*(?:일\|[0-9])", re.S)


	def _normalize(text, label, s, e):
	while s < e and text[s] in " .,\t\n": s += 1
	while e > s and text[e - 1] in " .,\t\n": e -= 1
	if label == "private_date":
	m = _DATE_END.match(text[s:e])
	if m and m.end() > 0:
	e = s + m.end()
	elif label in ("private_person", "personal_handle", "private_address"):
	for _ in range(2):
	seg = text[s:e]
	for j in _TRAILING_JOSA:
	if seg.endswith(j) and (e - s) - len(j) >= 2:
	e -= len(j); break
	else:
	break
	return s, e


	def extract_pii(text, max_length=256):
	enc = tokenizer(text, truncation=True, max_length=max_length,
	return_offsets_mapping=True, return_tensors="pt")
	offsets = enc.pop("offset_mapping")[0].tolist()
	with torch.no_grad():
	logits = model(**{k: v.to(model.device) for k, v in enc.items()}).logits
	pred = logits.argmax(-1)[0].tolist()
	id2label = model.config.id2label
	spans, active = [], None
	for i, lid in enumerate(pred):
	label = id2label[int(lid)]
	cs, ce = offsets[i]
	if cs == ce:
	if active: spans.append(active); active = None
	continue
	if label == "O":
	if active: spans.append(active); active = None
	continue
	prefix, cat = label.split("-", 1)
	if prefix in ("B", "S") or not active or active[0] != cat:
	if active: spans.append(active)
	active = [cat, cs, ce]
	else:
	active[2] = ce
	if active: spans.append(active)
	out = []
	for cat, s, e in spans:
	s, e = _normalize(text, cat, s, e)
	if text[s:e].strip():
	out.append({"label": cat, "start": s, "end": e, "text": text[s:e]})
	return out


	def redact(text):
	spans = sorted(extract_pii(text), key=lambda s: s["start"], reverse=True)
	for s in spans:
	text = text[:s["start"]] + f"[{s['label'].upper()}]" + text[s["end"]:]
	return text


	if __name__ == "__main__":
	for t in ["김민수님의 번호는 010-1234-5678입니다.",
	"계좌 110-234-567890으로 입금하고 minsu@example.com으로 알려주세요.",
	"이수진 고객님 생년월일은 1985년 3월 12일입니다."]:
	print(t)
	for sp in extract_pii(t):
	print(f" {sp['label']:16} {sp['text']!r}")
	print(" REDACT:", redact(t)); print()