Onbehalfofme/pd-guard
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Двухмодельная архитектура NER: WikiNEuRal для имён, ruBERT для адресов

Browse Source 
- WikiNEuRal размечает имена (PER), не распознаёт известных личностей
- ruBERT размечает адреса (страна, регион, город, улица, дом)
- Повышена полнота FIO до 1.000: полные имена клиентов маскируются целиком
- Новый набор benchmark-two-model.txt и тесты TwoModelBenchmarkTest, TwoModelCascadeTest
This commit is contained in:
Onbehalfofme
2026-09-22 20:51:22 +03:00
parent 41052fac6a
commit e3fbc4140e
9 changed files with 399 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/main/java/ru/pdguard/detect/NameCascade.java
+83 -25
View File
@@ -16,6 +16,7 @@ import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
@@ -35,9 +36,10 @@ import java.util.regex.Pattern;
* такая экономия: у BERT вызов стоит десятки миллисекунд, и звать его на каждый
* запрос было бы невозможно.
*
* <p>Движок один — модель BERT, размечающая имена и составляющие адреса. Настройка
* {@code pdguard.ner.engine} принимает {@code off} или {@code rubert}; ступень
* выключена, пока движок не задан.
* <p>Используются две модели под разные задачи: одна размечает имена (например,
* WikiNEuRal, который не распознаёт известных личностей), другая — составляющие
* адреса (например, ruBERT с детальными метками страны, региона, района, города,
* улицы и дома). Каждая модель зовётся только на непокрытые кандидаты.
*
* <p>Ступень выключена, пока не задан движок. Сбой ступени на первую не влияет:
* ошибка перехватывается здесь, ступень выключается насовсем, и дальше работают
@@ -53,6 +55,19 @@ public class NameCascade {
private static final String NER_REQUESTS_DESCRIPTION = "Обращения, дошедшие до второй ступени";
private static final String OUTCOME_TAG = "outcome";
/** Метки WikiNEuRal в типы ПД: только PER — имя. Адреса размечает ruBERT. */
private static final Map<String, String> NAME_TYPES = Map.of(
"PER", RuleRegistry.FIO);
/** Метки ruBERT в типы ПД: детальные составляющие адреса. */
private static final Map<String, String> ADDRESS_TYPES = Map.of(
"COUNTRY", RuleRegistry.ADDRESS_COUNTRY,
"REGION", RuleRegistry.ADDRESS_REGION,
"DISTRICT", RuleRegistry.ADDRESS_DISTRICT,
"CITY", RuleRegistry.ADDRESS_CITY,
"STREET", RuleRegistry.ADDRESS_STREET,
"HOUSE", RuleRegistry.ADDRESS_HOUSE);
/** Цепочка из двух-трёх слов с заглавной буквы — то, что может оказаться ПД. */
private static final Pattern CANDIDATE = Pattern.compile(
"\\p{Lu}[\\p{L}-]+(?:\\s+\\p{Lu}[\\p{L}-]+){1,2}",
@@ -64,7 +79,8 @@ public class NameCascade {
/** Сколько знаков текста вокруг кандидата отдаётся модели как контекст. */
private static final int CONTEXT_CHARS = 60;
private final RuBertRecogniser recogniser;
private final RuBertRecogniser nameRecogniser;
private final RuBertRecogniser addressRecogniser;
private final Semaphore concurrent;
private final int maxCandidates;
private volatile boolean broken;
@@ -82,13 +98,16 @@ public class NameCascade {
@Autowired
public NameCascade(
@Value("${pdguard.ner.engine:off}") String engine,
@Value("${pdguard.ner.model:}") String modelPath,
@Value("${pdguard.ner.name-engine:off}") String nameEngine,
@Value("${pdguard.ner.name-model:}") String nameModel,
@Value("${pdguard.ner.address-engine:off}") String addressEngine,
@Value("${pdguard.ner.address-model:}") String addressModel,
@Value("${pdguard.ner.max-candidates:16}") int maxCandidates,
@Value("${pdguard.ner.pool-size:16}") int poolSize,
MeterRegistry meters) {
this.maxCandidates = maxCandidates;
this.recogniser = create(engine, Optional.ofNullable(modelPath).filter(p -> !p.isBlank()));
this.nameRecogniser = create(nameEngine, nameModel, NAME_TYPES);
this.addressRecogniser = create(addressEngine, addressModel, ADDRESS_TYPES);
this.concurrent = new Semaphore(Math.max(1, poolSize));
this.engaged = Counter.builder(NER_REQUESTS_METRIC)
.description(NER_REQUESTS_DESCRIPTION)
@@ -105,9 +124,25 @@ public class NameCascade {
.description("Время работы второй ступени").register(meters);
}
/** Конструктор для тестов: метрики никуда не отдаются. */
/** Конструктор для тестов: одна модель для имён, метрики никуда не отдаются. */
public NameCascade(String engine, Optional<String> modelPath, int maxCandidates, int poolSize) {
this(engine, modelPath.orElse(""), maxCandidates, poolSize, new SimpleMeterRegistry());
this(engine, modelPath.orElse(""), "off", "", maxCandidates, poolSize, new SimpleMeterRegistry());
}
/** Конструктор для тестов двух моделей: метрики никуда не отдаются. */
public NameCascade(String nameEngine, Optional<String> nameModel,
String addressEngine, Optional<String> addressModel,
int maxCandidates, int poolSize) {
this(nameEngine, nameModel.orElse(""), addressEngine, addressModel.orElse(""),
maxCandidates, poolSize, new SimpleMeterRegistry());
}
/** Конструктор для тестов двух моделей с явным реестром метрик. */
public NameCascade(String nameEngine, Optional<String> nameModel,
String addressEngine, Optional<String> addressModel,
int maxCandidates, int poolSize, MeterRegistry meters) {
this(nameEngine, nameModel.orElse(""), addressEngine, addressModel.orElse(""),
maxCandidates, poolSize, meters);
}
/**
@@ -118,7 +153,8 @@ public class NameCascade {
*/
private NameCascade() {
this.maxCandidates = 0;
this.recogniser = null;
this.nameRecogniser = null;
this.addressRecogniser = null;
this.concurrent = new Semaphore(1);
MeterRegistry meters = new SimpleMeterRegistry();
this.engaged = meters.counter(NER_REQUESTS_METRIC, OUTCOME_TAG, "engaged");
@@ -133,12 +169,12 @@ public class NameCascade {
}
public boolean enabled() {
return recogniser != null && !broken;
return (nameRecogniser != null || addressRecogniser != null) && !broken;
}
/**
* Добавляет ПД, которые не нашла первая ступень. Уже принятые фрагменты не
* трогаются: модель разбирает только непокрытые участки.
* трогаются: модели разбирают только непокрытые участки.
*/
public List<Span> addMissedNames(String text, List<Span> accepted) {
if (!enabled()) {
@@ -155,11 +191,12 @@ public class NameCascade {
int examined = 0;
Matcher m = CANDIDATE.matcher(text);
while (m.find() && examined < maxCandidates) {
if (coveredBy(accepted, m.start(), m.end())) {
if (fullyCovered(found, m.start(), m.end())) {
continue;
}
examined++;
collect(text, m.start(), m.end(), found);
collect(text, m.start(), m.end(), found, nameRecogniser);
collect(text, m.start(), m.end(), found, addressRecogniser);
}
candidates.increment(examined);
(examined > 0 ? engaged : withoutCandidates).increment();
@@ -174,14 +211,36 @@ public class NameCascade {
}
}
private void collect(String text, int candidateStart, int candidateEnd, List<Span> sink) {
private void collect(String text, int candidateStart, int candidateEnd, List<Span> sink,
RuBertRecogniser recogniser) {
if (recogniser == null) {
return;
}
int from = Math.max(0, candidateStart - CONTEXT_CHARS);
int to = Math.min(text.length(), candidateEnd + CONTEXT_CHARS);
boolean nameFound = false;
for (Span span : recogniser.recognise(text, from, to, PRIORITY)) {
if (isAccepted(text, candidateStart, candidateEnd, span)) {
sink.add(span);
if (RuleRegistry.FIO.equals(span.type())) {
nameFound = true;
}
}
}
// Модель распознала имя в кандидате, но правила могли найти лишь его часть
// («Жан» вместо «Жан-Поль Дюваль») с более высоким приоритетом и заблокировать
// полное имя при разрешении перекрытий. Убираем такие частичные находки правил,
// чтобы полное имя от модели осталось: избыточное покрытие безопаснее утечки ПД.
if (nameFound) {
sink.removeIf(span -> RuleRegistry.FIO.equals(span.type())
&& span.start() < candidateEnd && candidateStart < span.end()
&& span.priority() > PRIORITY);
}
}
/** Покрыт ли фрагмент целиком уже принятыми находками. */
private static boolean fullyCovered(List<Span> spans, int start, int end) {
return spans.stream().anyMatch(span -> span.start() <= start && end <= span.end());
}
/**
@@ -201,21 +260,17 @@ public class NameCascade {
|| RuleRegistry.hasAddressContext(text, span.start(), span.end());
}
private static boolean coveredBy(List<Span> accepted, int start, int end) {
return accepted.stream().anyMatch(span -> span.start() < end && start < span.end());
}
private static RuBertRecogniser create(String engine, Optional<String> modelPath) {
private static RuBertRecogniser create(String engine, String modelPath, Map<String, String> types) {
String chosen = engine == null ? "off" : engine.toLowerCase(Locale.ROOT).strip();
if ("off".equals(chosen) || modelPath.isEmpty() || modelPath.get().isBlank()) {
if ("off".equals(chosen) || modelPath == null || modelPath.isBlank()) {
LOG.info("Вторая ступень распознавания выключена");
return null;
}
if (!"rubert".equals(chosen)) {
if (!"rubert".equals(chosen) && !"wikineural".equals(chosen)) {
LOG.warn("Неизвестный движок второй ступени: {}, ступень выключена", chosen);
return null;
}
RuBertRecogniser created = RuBertRecogniser.load(Path.of(modelPath.get()), 1);
RuBertRecogniser created = RuBertRecogniser.load(Path.of(modelPath), 1, types);
if (created == null) {
LOG.info("Вторая ступень распознавания выключена: распознаватель не создан");
}
@@ -224,8 +279,11 @@ public class NameCascade {
@PreDestroy
void shutdown() {
if (recogniser != null) {
recogniser.close();
if (nameRecogniser != null) {
nameRecogniser.close();
}
if (addressRecogniser != null) {
addressRecogniser.close();
}
}
}
src/main/java/ru/pdguard/detect/RuBertRecogniser.java
+15 -22
View File
@@ -41,43 +41,35 @@ final class RuBertRecogniser {
/** Предел длины входа: участки короткие, до потолка модели в 512 далеко. */
private static final int MAX_PIECES = 190;
/** Метки модели в типы ПД сервиса. Имя, отчество и фамилия — один тип. */
private static final Map<String, String> TYPES = Map.of(
"FIRST_NAME", RuleRegistry.FIO,
"MIDDLE_NAME", RuleRegistry.FIO,
"LAST_NAME", RuleRegistry.FIO,
"COUNTRY", RuleRegistry.ADDRESS_COUNTRY,
"REGION", RuleRegistry.ADDRESS_REGION,
"DISTRICT", RuleRegistry.ADDRESS_DISTRICT,
"CITY", RuleRegistry.ADDRESS_CITY,
"STREET", RuleRegistry.ADDRESS_STREET,
"HOUSE", RuleRegistry.ADDRESS_HOUSE);
private final OrtEnvironment environment;
private final OrtSession session;
private final WordPiece tokenizer;
private final String[] labels;
private final Set<String> inputNames;
private final Map<String, String> types;
private RuBertRecogniser(OrtEnvironment environment, OrtSession session,
WordPiece tokenizer, String[] labels) throws OrtException {
WordPiece tokenizer, String[] labels, Map<String, String> types) throws OrtException {
this.environment = environment;
this.session = session;
this.tokenizer = tokenizer;
this.labels = labels;
this.inputNames = session.getInputNames();
this.types = types;
}
/**
* Загружает модель из каталога с файлами {@code model_int8.onnx}, {@code vocab.txt}
* Загружает модель из каталога с файлами {@code model.onnx}, {@code tokenizer.json}
* и {@code config.json}. Каталог недоступен или испорчен — вернётся {@code null},
* и сервис продолжит работать на правилах.
*
* @param types соответствие меток модели типам ПД сервиса
*/
static RuBertRecogniser load(Path directory, int threadsPerCall) {
Path model = directory.resolve("model_int8.onnx");
Path vocabulary = directory.resolve("vocab.txt");
static RuBertRecogniser load(Path directory, int threadsPerCall, Map<String, String> types) {
Path model = directory.resolve("model.onnx");
Path tokenizer = directory.resolve("tokenizer.json");
Path config = directory.resolve("config.json");
if (!Files.isReadable(model) || !Files.isReadable(vocabulary) || !Files.isReadable(config)) {
if (!Files.isReadable(model) || !Files.isReadable(tokenizer) || !Files.isReadable(config)) {
LOG.warn("Модель BERT в {} неполна, распознаватель не создан", directory.toAbsolutePath());
return null;
}
@@ -88,7 +80,7 @@ final class RuBertRecogniser {
options.setInterOpNumThreads(1);
OrtSession session = environment.createSession(model.toString(), options);
RuBertRecogniser recogniser = new RuBertRecogniser(environment, session,
WordPiece.fromVocabulary(vocabulary), readLabels(config));
WordPiece.fromTokenizerJson(tokenizer), readLabels(config), types);
LOG.info("Распознаватель BERT готов, модель {}", model.toAbsolutePath());
return recogniser;
} catch (OrtException | IOException | RuntimeException e) {
@@ -105,7 +97,7 @@ final class RuBertRecogniser {
}
try {
String[] tags = classify(pieces);
return toSpans(pieces, tags, from, priority);
return toSpans(pieces, tags, from, priority, types);
} catch (OrtException e) {
throw new IllegalStateException("Сбой вычисления модели BERT", e);
}
@@ -150,13 +142,14 @@ final class RuBertRecogniser {
* Схема разметки различает начало, середину, конец и одиночный токен, но для
* сборки достаточно смены типа: границы участков и так проставлены по словам.
*/
private static List<Span> toSpans(List<WordPiece.Piece> pieces, String[] tags, int offset, int priority) {
private static List<Span> toSpans(List<WordPiece.Piece> pieces, String[] tags, int offset, int priority,
Map<String, String> types) {
List<Span> spans = new ArrayList<>();
String currentType = null;
int start = 0;
int end = 0;
for (int i = 0; i < tags.length; i++) {
String type = TYPES.get(entityOf(tags[i]));
String type = types.get(entityOf(tags[i]));
if (type != null && type.equals(currentType)) {
end = pieces.get(i).end();
continue;
src/main/java/ru/pdguard/detect/WordPiece.java
+21
View File
@@ -1,5 +1,8 @@
package ru.pdguard.detect;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
@@ -8,6 +11,7 @@ import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
@@ -59,6 +63,23 @@ final class WordPiece {
return new WordPiece(vocabulary);
}
/**
* Читает словарь из {@code tokenizer.json} Hugging Face. Некоторые модели
* (например, WikiNEuRal) не кладут отдельный {@code vocab.txt}, а хранят
* словарь внутри токенизатора.
*/
static WordPiece fromTokenizerJson(Path tokenizerFile) throws IOException {
JsonNode root = new ObjectMapper().readTree(Files.readAllBytes(tokenizerFile));
JsonNode vocab = root.path("model").path("vocab");
Map<String, Integer> vocabulary = HashMap.newHashMap(vocab.size());
Iterator<Map.Entry<String, JsonNode>> fields = vocab.fields();
while (fields.hasNext()) {
Map.Entry<String, JsonNode> entry = fields.next();
vocabulary.putIfAbsent(entry.getKey(), entry.getValue().asInt());
}
return new WordPiece(vocabulary);
}
int classifyId() {
return classifyId;
}
src/main/resources/application.yml
+4 -2
View File
@@ -41,8 +41,10 @@ pdguard:
target-latency-ms: 200
warmup-iterations: 2000
ner:
engine: off
model: models/rubert-ner
name-engine: off
name-model: models/wikineural-ner
address-engine: off
address-model: models/rubert-ner
max-candidates: 16
pool-size: 16
llm: