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

- WikiNEuRal размечает имена (PER), не распознаёт известных личностей
- ruBERT размечает адреса (страна, регион, город, улица, дом)
- Повышена полнота FIO до 1.000: полные имена клиентов маскируются целиком
- Новый набор benchmark-two-model.txt и тесты TwoModelBenchmarkTest, TwoModelCascadeTest

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 вызов стоит десятки миллисекунд, и звать его на каждый
  * запрос было бы невозможно.
  *
- * <p>Движок один — модель BERT, размечающая имена и составляющие адреса. Настройка
+ * <p>Используются две модели под разные задачи: одна размечает имена (например,
- * {@code pdguard.ner.engine} принимает {@code off} или {@code rubert}; ступень
+ * 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.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.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) {
+    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();
-        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) {
+        if (nameRecogniser != null) {
-            recogniser.close();
+            nameRecogniser.close();
+        }
+        if (addressRecogniser != null) {
+            addressRecogniser.close();
         }
     }
 }

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<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) {
+    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");
-        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

@@ -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

@@ -41,8 +41,10 @@ pdguard:
   target-latency-ms: 200
   warmup-iterations: 2000
   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
     pool-size: 16
   llm:

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++) {

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 для адресов.
+ *
+ * <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);
+    }
+}

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<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)),
+                "известная личность не должна маскироваться");
+    }
+}

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}}

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
+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"