pd-guard/src/test/java/ru/pdguard/TwoModelCascadeTest.java

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