Двухмодельная архитектура NER: WikiNEuRal для имён, ruBERT для адресов
- WikiNEuRal размечает имена (PER), не распознаёт известных личностей - ruBERT размечает адреса (страна, регион, город, улица, дом) - Повышена полнота FIO до 1.000: полные имена клиентов маскируются целиком - Новый набор benchmark-two-model.txt и тесты TwoModelBenchmarkTest, TwoModelCascadeTest
This commit is contained in:
@@ -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,15 +211,37 @@ 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();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -15,6 +15,7 @@ import java.nio.file.Path;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import java.util.List;
|
||||
import java.util.Optional;
|
||||
import java.util.concurrent.Callable;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.Executors;
|
||||
@@ -173,7 +174,7 @@ class PerformanceBenchmarkTest {
|
||||
MeterRegistry meters = new SimpleMeterRegistry();
|
||||
Pipeline withCascade = new Pipeline(new RuleRegistry(), new Masker(),
|
||||
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++) {
|
||||
|
||||
@@ -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
@@ -1,25 +1,45 @@
|
||||
#!/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
|
||||
set -euo pipefail
|
||||
cd "$(dirname "$0")/.."
|
||||
|
||||
BASE=https://huggingface.co/onnx-community/bert-base-NER-Russian-ONNX/resolve/main
|
||||
DIR=models/rubert-ner
|
||||
mkdir -p "$DIR"
|
||||
|
||||
echo "Загрузка модели в $DIR (около 170 МБ)"
|
||||
curl -sSL --progress-bar -o "$DIR/model_int8.onnx" "$BASE/onnx/model_int8.onnx"
|
||||
for f in vocab.txt config.json tokenizer.json tokenizer_config.json; do
|
||||
curl -sSL -o "$DIR/$f" "$BASE/$f"
|
||||
fetch() {
|
||||
local base="$1" dir="$2" model="$3"; shift 3
|
||||
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
|
||||
ls -lh "$DIR"
|
||||
echo
|
||||
echo "Готово. Включить: pdguard.ner.engine=rubert, pdguard.ner.model=$DIR"
|
||||
fetch "https://huggingface.co/rhnfzl/wikineural-multilingual-ner-onnx/resolve/main" \
|
||||
models/wikineural-ner model.onnx \
|
||||
config.json tokenizer.json tokenizer_config.json special_tokens_map.json
|
||||
|
||||
fetch "https://huggingface.co/onnx-community/bert-base-NER-Russian-ONNX/resolve/main" \
|
||||
models/rubert-ner onnx/model_int8.onnx \
|
||||
vocab.txt config.json tokenizer.json tokenizer_config.json
|
||||
|
||||
# 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"
|
||||
Reference in New Issue
Block a user