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

- 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
@@ -16,6 +16,7 @@ import java.nio.file.Path;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import java.util.Locale; import java.util.Locale;
import java.util.Map;
import java.util.Optional; import java.util.Optional;
import java.util.concurrent.Semaphore; import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeUnit;
@@ -35,9 +36,10 @@ import java.util.regex.Pattern;
* такая экономия: у BERT вызов стоит десятки миллисекунд, и звать его на каждый * такая экономия: у BERT вызов стоит десятки миллисекунд, и звать его на каждый
* запрос было бы невозможно. * запрос было бы невозможно.
* *
* <p>Движок один — модель BERT, размечающая имена и составляющие адреса. Настройка * <p>Используются две модели под разные задачи: одна размечает имена (например,
* {@code pdguard.ner.engine} принимает {@code off} или {@code rubert}; ступень * WikiNEuRal, который не распознаёт известных личностей), другая — составляющие
* выключена, пока движок не задан. * адреса (например, ruBERT с детальными метками страны, региона, района, города,
* улицы и дома). Каждая модель зовётся только на непокрытые кандидаты.
* *
* <p>Ступень выключена, пока не задан движок. Сбой ступени на первую не влияет: * <p>Ступень выключена, пока не задан движок. Сбой ступени на первую не влияет:
* ошибка перехватывается здесь, ступень выключается насовсем, и дальше работают * ошибка перехватывается здесь, ступень выключается насовсем, и дальше работают
@@ -53,6 +55,19 @@ public class NameCascade {
private static final String NER_REQUESTS_DESCRIPTION = "Обращения, дошедшие до второй ступени"; private static final String NER_REQUESTS_DESCRIPTION = "Обращения, дошедшие до второй ступени";
private static final String OUTCOME_TAG = "outcome"; 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( private static final Pattern CANDIDATE = Pattern.compile(
"\\p{Lu}[\\p{L}-]+(?:\\s+\\p{Lu}[\\p{L}-]+){1,2}", "\\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 static final int CONTEXT_CHARS = 60;
private final RuBertRecogniser recogniser; private final RuBertRecogniser nameRecogniser;
private final RuBertRecogniser addressRecogniser;
private final Semaphore concurrent; private final Semaphore concurrent;
private final int maxCandidates; private final int maxCandidates;
private volatile boolean broken; private volatile boolean broken;
@@ -82,13 +98,16 @@ public class NameCascade {
@Autowired @Autowired
public NameCascade( public NameCascade(
@Value("${pdguard.ner.engine:off}") String engine, @Value("${pdguard.ner.name-engine:off}") String nameEngine,
@Value("${pdguard.ner.model:}") String modelPath, @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.max-candidates:16}") int maxCandidates,
@Value("${pdguard.ner.pool-size:16}") int poolSize, @Value("${pdguard.ner.pool-size:16}") int poolSize,
MeterRegistry meters) { MeterRegistry meters) {
this.maxCandidates = maxCandidates; 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.concurrent = new Semaphore(Math.max(1, poolSize));
this.engaged = Counter.builder(NER_REQUESTS_METRIC) this.engaged = Counter.builder(NER_REQUESTS_METRIC)
.description(NER_REQUESTS_DESCRIPTION) .description(NER_REQUESTS_DESCRIPTION)
@@ -105,9 +124,25 @@ public class NameCascade {
.description("Время работы второй ступени").register(meters); .description("Время работы второй ступени").register(meters);
} }
/** Конструктор для тестов: метрики никуда не отдаются. */ /** Конструктор для тестов: одна модель для имён, метрики никуда не отдаются. */
public NameCascade(String engine, Optional<String> modelPath, int maxCandidates, int poolSize) { 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() { private NameCascade() {
this.maxCandidates = 0; this.maxCandidates = 0;
this.recogniser = null; this.nameRecogniser = null;
this.addressRecogniser = null;
this.concurrent = new Semaphore(1); this.concurrent = new Semaphore(1);
MeterRegistry meters = new SimpleMeterRegistry(); MeterRegistry meters = new SimpleMeterRegistry();
this.engaged = meters.counter(NER_REQUESTS_METRIC, OUTCOME_TAG, "engaged"); this.engaged = meters.counter(NER_REQUESTS_METRIC, OUTCOME_TAG, "engaged");
@@ -133,12 +169,12 @@ public class NameCascade {
} }
public boolean enabled() { public boolean enabled() {
return recogniser != null && !broken; return (nameRecogniser != null || addressRecogniser != null) && !broken;
} }
/** /**
* Добавляет ПД, которые не нашла первая ступень. Уже принятые фрагменты не * Добавляет ПД, которые не нашла первая ступень. Уже принятые фрагменты не
* трогаются: модель разбирает только непокрытые участки. * трогаются: модели разбирают только непокрытые участки.
*/ */
public List<Span> addMissedNames(String text, List<Span> accepted) { public List<Span> addMissedNames(String text, List<Span> accepted) {
if (!enabled()) { if (!enabled()) {
@@ -155,11 +191,12 @@ public class NameCascade {
int examined = 0; int examined = 0;
Matcher m = CANDIDATE.matcher(text); Matcher m = CANDIDATE.matcher(text);
while (m.find() && examined < maxCandidates) { while (m.find() && examined < maxCandidates) {
if (coveredBy(accepted, m.start(), m.end())) { if (fullyCovered(found, m.start(), m.end())) {
continue; continue;
} }
examined++; 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); candidates.increment(examined);
(examined > 0 ? engaged : withoutCandidates).increment(); (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 from = Math.max(0, candidateStart - CONTEXT_CHARS);
int to = Math.min(text.length(), candidateEnd + CONTEXT_CHARS); int to = Math.min(text.length(), candidateEnd + CONTEXT_CHARS);
boolean nameFound = false;
for (Span span : recogniser.recognise(text, from, to, PRIORITY)) { for (Span span : recogniser.recognise(text, from, to, PRIORITY)) {
if (isAccepted(text, candidateStart, candidateEnd, span)) { if (isAccepted(text, candidateStart, candidateEnd, span)) {
sink.add(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()); || RuleRegistry.hasAddressContext(text, span.start(), span.end());
} }
private static boolean coveredBy(List<Span> accepted, int start, int end) { private static RuBertRecogniser create(String engine, String modelPath, Map<String, String> types) {
return accepted.stream().anyMatch(span -> span.start() < end && start < span.end());
}
private static RuBertRecogniser create(String engine, Optional<String> modelPath) {
String chosen = engine == null ? "off" : engine.toLowerCase(Locale.ROOT).strip(); 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("Вторая ступень распознавания выключена"); LOG.info("Вторая ступень распознавания выключена");
return null; return null;
} }
if (!"rubert".equals(chosen)) { if (!"rubert".equals(chosen) && !"wikineural".equals(chosen)) {
LOG.warn("Неизвестный движок второй ступени: {}, ступень выключена", chosen); LOG.warn("Неизвестный движок второй ступени: {}, ступень выключена", chosen);
return null; return null;
} }
RuBertRecogniser created = RuBertRecogniser.load(Path.of(modelPath.get()), 1); RuBertRecogniser created = RuBertRecogniser.load(Path.of(modelPath), 1, types);
if (created == null) { if (created == null) {
LOG.info("Вторая ступень распознавания выключена: распознаватель не создан"); LOG.info("Вторая ступень распознавания выключена: распознаватель не создан");
} }
@@ -224,8 +279,11 @@ public class NameCascade {
@PreDestroy @PreDestroy
void shutdown() { void shutdown() {
if (recogniser != null) { if (nameRecogniser != null) {
recogniser.close(); nameRecogniser.close();
}
if (addressRecogniser != null) {
addressRecogniser.close();
} }
} }
} }
@@ -41,43 +41,35 @@ final class RuBertRecogniser {
/** Предел длины входа: участки короткие, до потолка модели в 512 далеко. */ /** Предел длины входа: участки короткие, до потолка модели в 512 далеко. */
private static final int MAX_PIECES = 190; 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 OrtEnvironment environment;
private final OrtSession session; private final OrtSession session;
private final WordPiece tokenizer; private final WordPiece tokenizer;
private final String[] labels; private final String[] labels;
private final Set<String> inputNames; private final Set<String> inputNames;
private final Map<String, String> types;
private RuBertRecogniser(OrtEnvironment environment, OrtSession session, 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.environment = environment;
this.session = session; this.session = session;
this.tokenizer = tokenizer; this.tokenizer = tokenizer;
this.labels = labels; this.labels = labels;
this.inputNames = session.getInputNames(); 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}, * и {@code config.json}. Каталог недоступен или испорчен — вернётся {@code null},
* и сервис продолжит работать на правилах. * и сервис продолжит работать на правилах.
*
* @param types соответствие меток модели типам ПД сервиса
*/ */
static RuBertRecogniser load(Path directory, int threadsPerCall) { static RuBertRecogniser load(Path directory, int threadsPerCall, Map<String, String> types) {
Path model = directory.resolve("model_int8.onnx"); Path model = directory.resolve("model.onnx");
Path vocabulary = directory.resolve("vocab.txt"); Path tokenizer = directory.resolve("tokenizer.json");
Path config = directory.resolve("config.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()); LOG.warn("Модель BERT в {} неполна, распознаватель не создан", directory.toAbsolutePath());
return null; return null;
} }
@@ -88,7 +80,7 @@ final class RuBertRecogniser {
options.setInterOpNumThreads(1); options.setInterOpNumThreads(1);
OrtSession session = environment.createSession(model.toString(), options); OrtSession session = environment.createSession(model.toString(), options);
RuBertRecogniser recogniser = new RuBertRecogniser(environment, session, RuBertRecogniser recogniser = new RuBertRecogniser(environment, session,
WordPiece.fromVocabulary(vocabulary), readLabels(config)); WordPiece.fromTokenizerJson(tokenizer), readLabels(config), types);
LOG.info("Распознаватель BERT готов, модель {}", model.toAbsolutePath()); LOG.info("Распознаватель BERT готов, модель {}", model.toAbsolutePath());
return recogniser; return recogniser;
} catch (OrtException | IOException | RuntimeException e) { } catch (OrtException | IOException | RuntimeException e) {
@@ -105,7 +97,7 @@ final class RuBertRecogniser {
} }
try { try {
String[] tags = classify(pieces); String[] tags = classify(pieces);
return toSpans(pieces, tags, from, priority); return toSpans(pieces, tags, from, priority, types);
} catch (OrtException e) { } catch (OrtException e) {
throw new IllegalStateException("Сбой вычисления модели BERT", 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<>(); List<Span> spans = new ArrayList<>();
String currentType = null; String currentType = null;
int start = 0; int start = 0;
int end = 0; int end = 0;
for (int i = 0; i < tags.length; i++) { 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)) { if (type != null && type.equals(currentType)) {
end = pieces.get(i).end(); end = pieces.get(i).end();
continue; continue;
@@ -1,5 +1,8 @@
package ru.pdguard.detect; package ru.pdguard.detect;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.BufferedReader; import java.io.BufferedReader;
import java.io.IOException; import java.io.IOException;
import java.io.InputStreamReader; import java.io.InputStreamReader;
@@ -8,6 +11,7 @@ import java.nio.file.Files;
import java.nio.file.Path; import java.nio.file.Path;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.HashMap; import java.util.HashMap;
import java.util.Iterator;
import java.util.List; import java.util.List;
import java.util.Map; import java.util.Map;
@@ -59,6 +63,23 @@ final class WordPiece {
return new WordPiece(vocabulary); 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() { int classifyId() {
return classifyId; return classifyId;
} }
+4 -2
View File
@@ -41,8 +41,10 @@ pdguard:
target-latency-ms: 200 target-latency-ms: 200
warmup-iterations: 2000 warmup-iterations: 2000
ner: ner:
engine: off name-engine: off
model: models/rubert-ner name-model: models/wikineural-ner
address-engine: off
address-model: models/rubert-ner
max-candidates: 16 max-candidates: 16
pool-size: 16 pool-size: 16
llm: llm:
@@ -15,6 +15,7 @@ import java.nio.file.Path;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.Arrays; import java.util.Arrays;
import java.util.List; import java.util.List;
import java.util.Optional;
import java.util.concurrent.Callable; import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService; import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors; import java.util.concurrent.Executors;
@@ -173,7 +174,7 @@ class PerformanceBenchmarkTest {
MeterRegistry meters = new SimpleMeterRegistry(); MeterRegistry meters = new SimpleMeterRegistry();
Pipeline withCascade = new Pipeline(new RuleRegistry(), new Masker(), Pipeline withCascade = new Pipeline(new RuleRegistry(), new Masker(),
new PayloadStore(10_000_000L, 30), new PayloadStore(10_000_000L, 30),
new NameCascade("rubert", model.toString(), 16, 4, meters)); new NameCascade("rubert", Optional.of(model.toString()), "off", Optional.empty(), 16, 4, meters));
// Прогрев второй ступени: модель инициализируется лениво, первые вызовы медленные. // Прогрев второй ступени: модель инициализируется лениво, первые вызовы медленные.
for (int i = 0; i < 200; i++) { for (int i = 0; i < 200; i++) {
@@ -0,0 +1,140 @@
package ru.pdguard;
import org.junit.jupiter.api.Test;
import ru.pdguard.config.SystemPolicy;
import ru.pdguard.core.PayloadStore;
import ru.pdguard.core.Pipeline;
import ru.pdguard.detect.NameCascade;
import ru.pdguard.detect.RuleRegistry;
import ru.pdguard.detect.Span;
import ru.pdguard.mask.Masker;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import static org.junit.jupiter.api.Assertions.assertTrue;
/**
* Оценка двухмодельной архитектуры: WikiNEuRal для имён, ruBERT для адресов.
*
* <p>Набор {@code benchmark-two-model.txt} проверяет, что имена клиентов и адреса
* маскируются, а известные личности — нет. Для каждого типа считается посимвольная
* точность, полнота и F1.
*/
class TwoModelBenchmarkTest {
private final Pipeline pipeline = new Pipeline(new RuleRegistry(), new Masker(),
new PayloadStore(10_000_000L, 30),
new NameCascade(
"wikineural", Optional.of("models/wikineural-ner"),
"rubert", Optional.of("models/rubert-ner"),
16, 4));
private static final class Score {
private int truePositive;
private int falsePositive;
private int falseNegative;
private int gold() {
return truePositive + falseNegative;
}
private double precision() {
int found = truePositive + falsePositive;
return found == 0 ? 1.0 : (double) truePositive / found;
}
private double recall() {
return gold() == 0 ? 1.0 : (double) truePositive / gold();
}
private double f1() {
double p = precision();
double r = recall();
return p + r == 0 ? 0.0 : 2 * p * r / (p + r);
}
}
@Test
void twoModelEfficiency() {
List<BenchmarkFixtures.Sample> samples = BenchmarkFixtures.load("/benchmark-two-model.txt");
Map<String, Score> byType = new LinkedHashMap<>();
Map<String, List<String>> missed = new LinkedHashMap<>();
for (BenchmarkFixtures.Sample sample : samples) {
List<Span> found = pipeline.findPersonalData(sample.text(), SystemPolicy.DEFAULT);
String[] goldChars = paint(sample.text().length(), sample.gold());
String[] foundChars = paint(sample.text().length(), found);
for (int i = 0; i < sample.text().length(); i++) {
account(byType, goldChars[i], foundChars[i]);
}
for (Span gold : sample.gold()) {
boolean hit = found.stream().anyMatch(f -> f.type().equals(gold.type()) && f.overlaps(gold));
if (!hit) {
missed.computeIfAbsent(gold.type(), t -> new ArrayList<>())
.add(sample.text().substring(gold.start(), gold.end()));
}
}
}
report(byType);
reportMissed(missed);
// Каждый тип должен быть найден с F1 не ниже 0.8.
for (Map.Entry<String, Score> e : byType.entrySet()) {
assertTrue(e.getValue().f1() >= 0.8,
String.format("F1 по типу %s упал до %.3f", e.getKey(), e.getValue().f1()));
}
}
private static String[] paint(int length, List<Span> spans) {
String[] painted = new String[length];
for (Span span : spans) {
for (int i = span.start(); i < Math.min(span.end(), length); i++) {
painted[i] = span.type();
}
}
return painted;
}
private static void account(Map<String, Score> byType, String gold, String found) {
if (gold != null) {
Score score = byType.computeIfAbsent(gold, t -> new Score());
if (gold.equals(found)) {
score.truePositive++;
} else {
score.falseNegative++;
}
}
if (found != null && !found.equals(gold)) {
byType.computeIfAbsent(found, t -> new Score()).falsePositive++;
}
}
private void report(Map<String, Score> byType) {
StringBuilder out = new StringBuilder(2048);
out.append("\n=== Эффективность двухмодельной архитектуры ===\n\n");
out.append(String.format("%-20s %8s %8s %8s %8s%n", "тип", "знаков", "точность", "полнота", "F1"));
byType.entrySet().stream()
.sorted(Comparator.comparingInt((Map.Entry<String, Score> e) -> e.getValue().gold()).reversed())
.forEach(e -> out.append(String.format("%-20s %8d %8.3f %8.3f %8.3f%n",
e.getKey(), e.getValue().gold(), e.getValue().precision(),
e.getValue().recall(), e.getValue().f1())));
System.out.println(out);
}
private void reportMissed(Map<String, List<String>> missed) {
if (missed.isEmpty()) {
return;
}
StringBuilder out = new StringBuilder();
out.append("\n=== Не распознанные значения по типам ===\n");
missed.forEach((type, values) -> out.append(type).append(": ")
.append(String.join(" | ", values)).append('\n'));
System.out.println(out);
}
}
@@ -0,0 +1,51 @@
package ru.pdguard;
import org.junit.jupiter.api.Test;
import ru.pdguard.config.SystemPolicy;
import ru.pdguard.core.PayloadStore;
import ru.pdguard.core.Pipeline;
import ru.pdguard.detect.NameCascade;
import ru.pdguard.detect.RuleRegistry;
import ru.pdguard.detect.Span;
import ru.pdguard.mask.Masker;
import java.util.List;
import java.util.Optional;
import static org.junit.jupiter.api.Assertions.assertTrue;
/** Две модели: WikiNEuRal для имён, ruBERT для адресов. */
class TwoModelCascadeTest {
private List<Span> find(String text) {
NameCascade cascade = new NameCascade(
"wikineural", Optional.of("models/wikineural-ner"),
"rubert", Optional.of("models/rubert-ner"),
16, 4);
Pipeline p = new Pipeline(new RuleRegistry(), new Masker(), new PayloadStore(1_000_000L, 30), cascade);
return p.findPersonalData(text, SystemPolicy.DEFAULT);
}
@Test
void recognisesNamesAndAddresses() {
String text = "Клиент Иванов Иван Иванович, проживает в городе Москва, на улице Тверская";
List<Span> spans = find(text);
System.out.println("TEXT: " + text);
for (Span s : spans) {
System.out.println(" -> " + s.type() + " [" + text.substring(s.start(), s.end()) + "]");
}
assertTrue(spans.stream().anyMatch(s -> s.type().equals(RuleRegistry.FIO)), "должно найти ФИО");
}
@Test
void wellKnownNamesNotMasked() {
String text = "Напиши стихотворение в духе Александра Пушкина про осень";
List<Span> spans = find(text);
System.out.println("TEXT: " + text);
for (Span s : spans) {
System.out.println(" -> " + s.type() + " [" + text.substring(s.start(), s.end()) + "]");
}
assertTrue(spans.stream().noneMatch(s -> s.type().equals(RuleRegistry.FIO)),
"известная личность не должна маскироваться");
}
}
@@ -0,0 +1,47 @@
# НАБОР ДЛЯ ДВУХМОДЕЛЬНОЙ АРХИТЕКТУРЫ. Проверяет вторую ступень: WikiNEuRal
# размечает имена (включая иностранные и нестандартные), ruBERT — составляющие
# адреса. Ключевая проверка — известные личности не маскируются, а имена клиентов
# и адреса маскируются. Формат: {{ТИП:значение}}.
# --- Имена клиентов (должны маскироваться) ---
Клиент {{FIO:Иванов Иван Иванович}} пришёл в офис
Заявление подал {{FIO:Смирнов Артём Викторович}}
Документы подписал {{FIO:Соколов А.П.}}
Клиент {{FIO:Дмитрий Волков}} закрыл счёт
Заёмщик {{FIO:Орлова}} внесла платёж
Поручитель {{FIO:Егоров}} не ответил на звонок
# --- Иностранные и нестандартные имена (WikiNEuRal должен распознать) ---
Клиент {{FIO:Жан-Поль Дюваль}} оформил заявку
Заявление подал {{FIO:Виктор Моро}} из Франции
Клиент {{FIO:Оноре де Роше}} открыл счёт
Документы подписал {{FIO:Готье Руссо}}
# --- Известные личности (НЕ должны маскироваться) ---
Напиши стихотворение в духе Александра Пушкина про осень
Портрет Толстого висит в кабинете
Мы изучаем сонеты Шекспира на этой неделе
Лев Толстой написал роман «Война и мир»
Собрание сочинений Достоевского стоит на полке
# --- Адреса (должны маскироваться) ---
Клиент проживает в г. {{ADDRESS_CITY:Москва}}, на улице {{ADDRESS_STREET:Тверская}}
Адрес: {{ADDRESS_POSTCODE:125009}}, г. {{ADDRESS_CITY:Москва}}, ул. {{ADDRESS_STREET:Тверская}}, д. {{ADDRESS_HOUSE:7}}, кв. {{ADDRESS_FLAT:15}}
Регион {{ADDRESS_REGION:Татарстан}}, город {{ADDRESS_CITY:Подольск}}
Страна проживания {{ADDRESS_COUNTRY:Россия}}, город {{ADDRESS_CITY:Санкт-Петербург}}
# --- Нестандартные адреса (ruBERT должен распознать) ---
Клиент проживает в г. {{ADDRESS_CITY:Новосибирск}}, на проспекте {{ADDRESS_STREET:Ленина}}
Адрес регистрации: г. {{ADDRESS_CITY:Екатеринбург}}, улица {{ADDRESS_STREET:Малышева}}, дом {{ADDRESS_HOUSE:51}}
# --- Смешанные случаи: имя + адрес ---
Клиент {{FIO:Петров Пётр}}, проживает в г. {{ADDRESS_CITY:Казань}}, на улице {{ADDRESS_STREET:Баумана}}
Заявление от {{FIO:Сидорова Анна}}, адрес: г. {{ADDRESS_CITY:Самара}}, ул. {{ADDRESS_STREET:Ленинградская}}, д. {{ADDRESS_HOUSE:10}}
# --- Ложные срабатывания: адрес отделения банка (НЕ маскируется) ---
Отделение банка находится по адресу: г. Москва, ул. Тверская, д. 7
Банкомат установлен в отделении на улице Арбат, дом 12
# --- Имя с ПД другого типа (однофамилец известной личности маскируется) ---
Клиент Александр Пушкин, паспорт {{PASSPORT:4509 123456}}
Клиент Лев Толстой, телефон {{PHONE:+7 916 123-45-67}}
+36 -16
View File
@@ -1,25 +1,45 @@
#!/usr/bin/env bash #!/usr/bin/env bash
# Загрузка модели распознавания для второй ступени. # Загрузка моделей распознавания для второй ступени.
# #
# Модель — BERT (onnx-community/bert-base-NER-Russian-ONNX, лицензия MIT): размечает # Используются две модели под разные задачи:
# имена и составляющие адреса вплоть до региона и района. Около 170 МБ, в репозиторий # - WikiNEuRal (Babelscape/wikineural-multilingual-ner, CC BY-NC-SA 4.0) — имена.
# не кладётся. # Многоязычный NER на базе mBERT, обучен на 9 языках, включая русский. Не
# распознаёт известных личностей («Лев Толстой»), поэтому справочник имён
# не нужен. ONNX-версия около 700 МБ.
# - ruBERT (onnx-community/bert-base-NER-Russian-ONNX, MIT) — адреса. Размечает
# страну, регион, район, город, улицу и дом отдельно. Около 170 МБ.
#
# Модели в репозиторий не кладутся.
# #
# ./tools/fetch-ner-model.sh # ./tools/fetch-ner-model.sh
set -euo pipefail set -euo pipefail
cd "$(dirname "$0")/.." cd "$(dirname "$0")/.."
BASE=https://huggingface.co/onnx-community/bert-base-NER-Russian-ONNX/resolve/main fetch() {
DIR=models/rubert-ner local base="$1" dir="$2" model="$3"; shift 3
mkdir -p "$DIR" mkdir -p "$dir"
echo "Загрузка модели в $dir"
curl -sSL --progress-bar -o "$dir/$model" "$base/$model"
for f in "$@"; do
curl -sSL -o "$dir/$f" "$base/$f"
done
echo
ls -lh "$dir"
echo
}
echo "Загрузка модели в $DIR (около 170 МБ)" fetch "https://huggingface.co/rhnfzl/wikineural-multilingual-ner-onnx/resolve/main" \
curl -sSL --progress-bar -o "$DIR/model_int8.onnx" "$BASE/onnx/model_int8.onnx" models/wikineural-ner model.onnx \
for f in vocab.txt config.json tokenizer.json tokenizer_config.json; do config.json tokenizer.json tokenizer_config.json special_tokens_map.json
curl -sSL -o "$DIR/$f" "$BASE/$f"
done
echo fetch "https://huggingface.co/onnx-community/bert-base-NER-Russian-ONNX/resolve/main" \
ls -lh "$DIR" models/rubert-ner onnx/model_int8.onnx \
echo vocab.txt config.json tokenizer.json tokenizer_config.json
echo "Готово. Включить: pdguard.ner.engine=rubert, pdguard.ner.model=$DIR"
# ruBERT кладёт модель в подкаталог onnx/; приведём к ожидаемому имени.
mv -f models/rubert-ner/onnx/model_int8.onnx models/rubert-ner/model.onnx 2>/dev/null || true
rmdir models/rubert-ner/onnx 2>/dev/null || true
echo "Готово. Включить:"
echo " pdguard.ner.name-engine=wikineural, pdguard.ner.name-model=models/wikineural-ner"
echo " pdguard.ner.address-engine=rubert, pdguard.ner.address-model=models/rubert-ner"