diff --git a/src/main/java/ru/pdguard/detect/NameCascade.java b/src/main/java/ru/pdguard/detect/NameCascade.java index c8b1056..735b129 100644 --- a/src/main/java/ru/pdguard/detect/NameCascade.java +++ b/src/main/java/ru/pdguard/detect/NameCascade.java @@ -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 вызов стоит десятки миллисекунд, и звать его на каждый * запрос было бы невозможно. * - *

Движок один — модель BERT, размечающая имена и составляющие адреса. Настройка - * {@code pdguard.ner.engine} принимает {@code off} или {@code rubert}; ступень - * выключена, пока движок не задан. + *

Используются две модели под разные задачи: одна размечает имена (например, + * WikiNEuRal, который не распознаёт известных личностей), другая — составляющие + * адреса (например, ruBERT с детальными метками страны, региона, района, города, + * улицы и дома). Каждая модель зовётся только на непокрытые кандидаты. * *

Ступень выключена, пока не задан движок. Сбой ступени на первую не влияет: * ошибка перехватывается здесь, ступень выключается насовсем, и дальше работают @@ -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 NAME_TYPES = Map.of( + "PER", RuleRegistry.FIO); + + /** Метки ruBERT в типы ПД: детальные составляющие адреса. */ + private static final Map 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 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 nameModel, + String addressEngine, Optional addressModel, + int maxCandidates, int poolSize) { + this(nameEngine, nameModel.orElse(""), addressEngine, addressModel.orElse(""), + maxCandidates, poolSize, new SimpleMeterRegistry()); + } + + /** Конструктор для тестов двух моделей с явным реестром метрик. */ + public NameCascade(String nameEngine, Optional nameModel, + String addressEngine, Optional 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 addMissedNames(String text, List 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 sink) { +private void collect(String text, int candidateStart, int candidateEnd, List 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 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 accepted, int start, int end) { - return accepted.stream().anyMatch(span -> span.start() < end && start < span.end()); - } - - private static RuBertRecogniser create(String engine, Optional modelPath) { + private static RuBertRecogniser create(String engine, String modelPath, Map 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(); } } } \ No newline at end of file diff --git a/src/main/java/ru/pdguard/detect/RuBertRecogniser.java b/src/main/java/ru/pdguard/detect/RuBertRecogniser.java index a04bba4..a24b747 100644 --- a/src/main/java/ru/pdguard/detect/RuBertRecogniser.java +++ b/src/main/java/ru/pdguard/detect/RuBertRecogniser.java @@ -41,43 +41,35 @@ final class RuBertRecogniser { /** Предел длины входа: участки короткие, до потолка модели в 512 далеко. */ private static final int MAX_PIECES = 190; - /** Метки модели в типы ПД сервиса. Имя, отчество и фамилия — один тип. */ - private static final Map 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 inputNames; + private final Map types; private RuBertRecogniser(OrtEnvironment environment, OrtSession session, - WordPiece tokenizer, String[] labels) throws OrtException { + WordPiece tokenizer, String[] labels, Map 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 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 toSpans(List pieces, String[] tags, int offset, int priority) { + private static List toSpans(List pieces, String[] tags, int offset, int priority, + Map types) { List 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; diff --git a/src/main/java/ru/pdguard/detect/WordPiece.java b/src/main/java/ru/pdguard/detect/WordPiece.java index 16f7163..6074e88 100644 --- a/src/main/java/ru/pdguard/detect/WordPiece.java +++ b/src/main/java/ru/pdguard/detect/WordPiece.java @@ -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 vocabulary = HashMap.newHashMap(vocab.size()); + Iterator> fields = vocab.fields(); + while (fields.hasNext()) { + Map.Entry entry = fields.next(); + vocabulary.putIfAbsent(entry.getKey(), entry.getValue().asInt()); + } + return new WordPiece(vocabulary); + } + int classifyId() { return classifyId; } diff --git a/src/main/resources/application.yml b/src/main/resources/application.yml index c5cd9de..fdada25 100644 --- a/src/main/resources/application.yml +++ b/src/main/resources/application.yml @@ -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: diff --git a/src/test/java/ru/pdguard/PerformanceBenchmarkTest.java b/src/test/java/ru/pdguard/PerformanceBenchmarkTest.java index ed4b30e..dddbf33 100644 --- a/src/test/java/ru/pdguard/PerformanceBenchmarkTest.java +++ b/src/test/java/ru/pdguard/PerformanceBenchmarkTest.java @@ -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++) { diff --git a/src/test/java/ru/pdguard/TwoModelBenchmarkTest.java b/src/test/java/ru/pdguard/TwoModelBenchmarkTest.java new file mode 100644 index 0000000..511f635 --- /dev/null +++ b/src/test/java/ru/pdguard/TwoModelBenchmarkTest.java @@ -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 для адресов. + * + *

Набор {@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 samples = BenchmarkFixtures.load("/benchmark-two-model.txt"); + Map byType = new LinkedHashMap<>(); + Map> missed = new LinkedHashMap<>(); + + for (BenchmarkFixtures.Sample sample : samples) { + List 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 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 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 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 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 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> 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); + } +} \ No newline at end of file diff --git a/src/test/java/ru/pdguard/TwoModelCascadeTest.java b/src/test/java/ru/pdguard/TwoModelCascadeTest.java new file mode 100644 index 0000000..389386c --- /dev/null +++ b/src/test/java/ru/pdguard/TwoModelCascadeTest.java @@ -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 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 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 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)), + "известная личность не должна маскироваться"); + } +} diff --git a/src/test/resources/benchmark-two-model.txt b/src/test/resources/benchmark-two-model.txt new file mode 100644 index 0000000..765cf8f --- /dev/null +++ b/src/test/resources/benchmark-two-model.txt @@ -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}} \ No newline at end of file diff --git a/tools/fetch-ner-model.sh b/tools/fetch-ner-model.sh index e8aa3e4..b2a7d1e 100755 --- a/tools/fetch-ner-model.sh +++ b/tools/fetch-ner-model.sh @@ -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" +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 "Загрузка модели в $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" -done +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 -echo -ls -lh "$DIR" -echo -echo "Готово. Включить: pdguard.ner.engine=rubert, pdguard.ner.model=$DIR" +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" \ No newline at end of file