RepinoidandCopilot
1b0bcd298a
Expose extraction, analysis, and replacement progress in the file table; preserve row mapping for duplicates and collapsed archives. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
347 lines
16 KiB
Python
347 lines
16 KiB
Python
"""
|
||
Проход 1: обнаружение сущностей в тексте документов.
|
||
|
||
Два метода сканирования:
|
||
1. scan_regex — быстрое regex-обнаружение (телефоны, email, ИНН, счета, компании, ФИО)
|
||
2. scan_llm_ner — универсальное LLM-обнаружение (любые приватные данные)
|
||
|
||
Замена: гибридные токены — читаемый шаблон + уникальный номер.
|
||
Пример: Иванов_0001, ООО_Технология_0034, ул_Ленина_0015
|
||
"""
|
||
|
||
import re
|
||
import json
|
||
import random
|
||
import time
|
||
import logging
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from typing import Dict, List, Optional
|
||
|
||
from .config import (
|
||
ENTITY_PATTERNS, COMPANY_PATTERN, PERSON_PATTERN,
|
||
RU_SURNAMES, COMPANY_NOUNS, CITY_POOL, STREET_POOL,
|
||
)
|
||
|
||
log = logging.getLogger("drhider")
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
# Гибридные токены: тип → (пул_шаблонов, префикс_для_счётчика)
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
TYPE_POOLS = {
|
||
"person_name": (RU_SURNAMES, "person"),
|
||
"person": (RU_SURNAMES, "person"),
|
||
"company": (COMPANY_NOUNS, "company"),
|
||
"phone": (["+7_000_000"], "phone"),
|
||
"email": (["email"], "email"),
|
||
"address": (CITY_POOL, "address"),
|
||
"inn": (["ИНН"], "inn"),
|
||
"inn_ul": (["ИНН"], "inn"),
|
||
"inn_fl": (["ИНН"], "inn"),
|
||
"kpp": (["КПП"], "kpp"),
|
||
"ogrn": (["ОГРН"], "ogrn"),
|
||
"bik": (["БИК"], "bik"),
|
||
"passport": (["Паспорт"], "passport"),
|
||
"bank_account": (["Счёт"], "bank"),
|
||
"rs": (["РС"], "bank"),
|
||
"ks": (["КС"], "bank"),
|
||
"contract_number": (["Договор"], "contract"),
|
||
"other_id": (["ID"], "other"),
|
||
}
|
||
_FALLBACK_POOL = (["Данные"], "data")
|
||
|
||
|
||
def _next_token(entity_type: str, counters: Dict[str, int]) -> str:
|
||
"""Сгенерировать токен замены: случайный шаблон из пула + номер.
|
||
|
||
Одна и та же сущность в mapping получает один и тот же токен
|
||
(mapping dict гарантирует согласованность).
|
||
Разные сущности одного типа получают РАЗНЫЕ шаблоны из пула.
|
||
|
||
Args:
|
||
entity_type: Тип сущности (person_name, company, phone...)
|
||
counters: Глобальные счётчики {prefix_key: count}
|
||
|
||
Returns:
|
||
Строка вида "Иванов_0001", "ООО_Технология_0034", "[Адрес_0015]"
|
||
"""
|
||
pool, prefix_key = TYPE_POOLS.get(entity_type, _FALLBACK_POOL)
|
||
template = random.choice(pool)
|
||
key = f"{prefix_key}_{template}"
|
||
counters[key] = counters.get(key, 0) + 1
|
||
return f"{template}_{counters[key]:04d}"
|
||
|
||
log = logging.getLogger("drhider")
|
||
|
||
|
||
def scan_regex(text: str, mapping: Dict[str, str], counters: Dict[str, int]) -> None:
|
||
"""Сканировать текст regex-паттернами, заполнить словарь замен.
|
||
|
||
Замена: гибридные токены (шаблон + номер).
|
||
|
||
Args:
|
||
text: Текст документа для сканирования
|
||
mapping: Словарь замен (мутабельный, пополняется)
|
||
counters: Глобальные счётчики токенов (мутабельный)
|
||
"""
|
||
# ── Сущности по regex-паттернам ──
|
||
for entity_type, pattern in ENTITY_PATTERNS.items():
|
||
for match in re.finditer(pattern, text, re.IGNORECASE | re.MULTILINE):
|
||
original = match.group(0).strip()
|
||
if not original or original in mapping:
|
||
continue
|
||
mapping[original] = _next_token(entity_type, counters)
|
||
|
||
# ── Названия компаний ──
|
||
for match in COMPANY_PATTERN.finditer(text):
|
||
original = match.group(0).strip()
|
||
if not original or original in mapping:
|
||
continue
|
||
mapping[original] = _next_token("company", counters)
|
||
|
||
# ── ФИО ──
|
||
for match in PERSON_PATTERN.finditer(text):
|
||
original = match.group(0).strip()
|
||
if not original or original in mapping:
|
||
continue
|
||
mapping[original] = _next_token("person_name", counters)
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
# LLM-NER: пофайловый чанкинг (Sonnet-схема)
|
||
# ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
class CancelRequested(Exception):
|
||
"""Обработка прервана пользователем (мягкая остановка между файлами LLM)."""
|
||
|
||
def __init__(self, llm_failures: int = 0):
|
||
self.llm_failures = llm_failures
|
||
|
||
# Приоритет границ для чанков (от предпочтительных к жёстким)
|
||
_CHUNK_BOUNDARIES = ['\n\n', '\n', '. ', '? ', '! ', '; ', ', ', ' ']
|
||
_CHUNK_SIZE = 6000 # символов на чанк (≈1500-2000 токенов RU — NER-качество)
|
||
_CHUNK_OVERLAP = 500 # перекрытие чанков (сущности на стыке)
|
||
_LLM_CONCURRENCY = 1 # параллельность убрана — всё последовательно (liberta/LLM не тянет конкурентность)
|
||
|
||
|
||
def split_into_chunks(text: str, size: int = _CHUNK_SIZE,
|
||
overlap: int = _CHUNK_OVERLAP) -> List[str]:
|
||
"""Разбить текст на чанки ~size символов с перекрытием overlap.
|
||
|
||
Режем по границам (\n\n, \n, ..., пробел) — чтобы не разрывать слова/сущности.
|
||
"""
|
||
chunks: List[str] = []
|
||
start = 0
|
||
tl = len(text)
|
||
while start < tl:
|
||
end = min(start + size, tl)
|
||
if end == tl:
|
||
chunks.append(text[start:])
|
||
break
|
||
cut = None
|
||
for boundary in _CHUNK_BOUNDARIES:
|
||
pos = text.rfind(boundary, start + size // 2, end)
|
||
if pos != -1:
|
||
cut = pos + len(boundary)
|
||
break
|
||
if cut is None:
|
||
cut = end # жёсткий разрез — нет ни одного пробела/границы
|
||
chunks.append(text[start:cut])
|
||
# overlap: отступить назад, найти начало слова
|
||
ov_start = max(start, cut - overlap)
|
||
space = text.find(' ', ov_start)
|
||
start = (space + 1) if (space != -1 and space < cut) else ov_start
|
||
return chunks
|
||
|
||
|
||
def normalize_entity(s: str) -> str:
|
||
"""Нормализовать строку сущности для надёжного дедупа (регистр/пробелы/тире)."""
|
||
s = s.strip()
|
||
s = re.sub(r'\s+', ' ', s)
|
||
s = s.lower()
|
||
s = re.sub(r'[«»\u201c\u201d\u2018\u2019"\u0027]', '"', s)
|
||
s = re.sub(r'[\u2014\u2013\u2212-]', '-', s)
|
||
s = s.replace('\u00ad', '') # мягкий перенос
|
||
s = s.replace('ё', 'е') # Ё→Е (OCR/PDF)
|
||
return s
|
||
|
||
|
||
def _verify_entity(value: str, full_text: str) -> bool:
|
||
"""Верификация: сущность реально есть в полном тексте файла (отсекает галлюцинации)."""
|
||
if value in full_text:
|
||
return True
|
||
return normalize_entity(value) in normalize_entity(full_text)
|
||
|
||
|
||
def _build_llm_prompt(text: str, already_found: List[str]) -> str:
|
||
"""Собрать промпт для одного текстового фрагмента (чанк/файл)."""
|
||
af = "\n".join(f"- {k}" for k in already_found) if already_found else "(none)"
|
||
return (
|
||
"You are a PII (Personally Identifiable Information) detector for Russian legal documents.\n\n"
|
||
"Find ALL sensitive or private information that is NOT already captured below.\n\n"
|
||
|
||
"Already captured by regex (skip these exact strings):\n"
|
||
f"{af}\n\n"
|
||
|
||
"Return a JSON array. Each item must have:\n"
|
||
' - "type": MUST be one of: person_name, company, phone, email, address,\n'
|
||
' inn, kpp, ogrn, bik, passport, bank_account, contract_number, other_id\n'
|
||
' Use "other_id" only if nothing else fits. If unsure — do NOT include.\n'
|
||
' - "value": the EXACT string from the text (copy verbatim)\n\n'
|
||
|
||
"NOT PII — do NOT include these (legal terms, not private data):\n"
|
||
"- Roles: Исполнитель, Заказчик, Стороны, Сторона, Покупатель, Поставщик\n"
|
||
"- Titles: Генеральный директор, Директор, действующего на основании\n"
|
||
"- Legal terms: Услуги, Договор, Приложение, Спецификация, НДС, Устава\n\n"
|
||
|
||
"Rules:\n"
|
||
"- Copy value VERBATIM — same spaces, punctuation, case as in text\n"
|
||
"- Same value → include only once\n"
|
||
"- Return ONLY the JSON array, no other text, no markdown wrapping\n\n"
|
||
|
||
f"Text:\n{text}"
|
||
)
|
||
|
||
|
||
def _call_llm(text: str, llm_client) -> List[dict]:
|
||
"""Один вызов LLM по фрагменту → проверенный список сущностей."""
|
||
raw = llm_client.complete(text).strip()
|
||
if raw.startswith("```"):
|
||
lines = raw.splitlines()
|
||
if lines and lines[0].startswith("```"):
|
||
lines = lines[1:]
|
||
if lines and lines[-1].strip() == "```":
|
||
lines = lines[:-1]
|
||
raw = "\n".join(lines).strip()
|
||
|
||
try:
|
||
entities = json.loads(raw)
|
||
except json.JSONDecodeError as original_error:
|
||
# Модель иногда добавляет текст вокруг массива; raw_decode корректно
|
||
# учитывает вложенные массивы и квадратные скобки внутри JSON-строк.
|
||
decoder = json.JSONDecoder()
|
||
entities = None
|
||
for index, char in enumerate(raw):
|
||
if char != "[":
|
||
continue
|
||
try:
|
||
candidate, _ = decoder.raw_decode(raw, index)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
if isinstance(candidate, list):
|
||
entities = candidate
|
||
break
|
||
if entities is None:
|
||
raise ValueError("LLM response does not contain a JSON array") from original_error
|
||
|
||
if not isinstance(entities, list):
|
||
raise ValueError("LLM response must be a JSON array")
|
||
for entity in entities:
|
||
if not isinstance(entity, dict):
|
||
raise ValueError("LLM entity must be a JSON object")
|
||
value = entity.get("value")
|
||
if not isinstance(value, str) or not value.strip():
|
||
raise ValueError("LLM entity value must be a non-empty string")
|
||
entity_type = entity.get("type")
|
||
if entity_type is not None and not isinstance(entity_type, str):
|
||
raise ValueError("LLM entity type must be a string")
|
||
return entities
|
||
|
||
|
||
def scan_llm_ner(all_texts: Dict[str, str], mapping: Dict[str, str],
|
||
llm_client, counters: Dict[str, int],
|
||
cancel_event=None, file_progress=None) -> int:
|
||
"""Универсальное LLM-обнаружение приватных данных (пофайлово, целиком).
|
||
|
||
Каждый файл обрабатывается ПОЛНОСТЬЮ: текст разбивается на чанки
|
||
(_CHUNK_SIZE=6000, overlap=_CHUNK_OVERLAP), чанки обрабатываются последовательно
|
||
(_LLM_CONCURRENCY), сущности дедуплицируются и верифицируются против полного
|
||
текста файла (отсечка галлюцинаций). regex-найденное не дублируется.
|
||
|
||
Args:
|
||
all_texts: {filename: text_content} — тексты всех файлов
|
||
mapping: Словарь замен (мутабельный, пополняется)
|
||
llm_client: Объект с методом .complete(prompt) -> str
|
||
counters: Глобальные счётчики токенов (мутабельный)
|
||
cancel_event: threading.Event — мягкая остановка между файлами (текущий добирается до конца)
|
||
file_progress: callable(event, fname, **fields) — прогресс по файлам/чанкам
|
||
"""
|
||
already_found = list(mapping.keys())
|
||
|
||
# Порядок файлов: крупные вперёд (приоритет на долгую обработку)
|
||
file_items = sorted(
|
||
all_texts.items(), key=lambda kv: len(kv[1]), reverse=True
|
||
)
|
||
llm_failures = 0
|
||
|
||
for fname, full_text in file_items:
|
||
# Мягкая остановка ТОЛЬКО между файлами — текущий файл добирается до конца
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
raise CancelRequested(llm_failures)
|
||
|
||
if not full_text or full_text.startswith("[DOC binary"):
|
||
continue
|
||
chunks = split_into_chunks(full_text)
|
||
|
||
if file_progress:
|
||
file_progress("file_start", fname, chars=len(full_text), chunks=len(chunks))
|
||
|
||
t0 = time.time()
|
||
|
||
# Параллельные вызовы LLM по чанкам этого файла
|
||
if len(chunks) > 1 and _LLM_CONCURRENCY > 1:
|
||
with ThreadPoolExecutor(max_workers=_LLM_CONCURRENCY) as _ex:
|
||
futures = [_ex.submit(_call_llm, _build_llm_prompt(c, mapping.keys()), llm_client)
|
||
for c in chunks]
|
||
chunk_results = [None] * len(futures)
|
||
future_indexes = {future: index for index, future in enumerate(futures)}
|
||
completed = 0
|
||
for fut in as_completed(futures):
|
||
index = future_indexes[fut]
|
||
try:
|
||
chunk_results[index] = fut.result()
|
||
except Exception as e:
|
||
llm_failures += 1
|
||
log.warning("LLM NER chunk failed in %r: %s", fname, e)
|
||
chunk_results[index] = []
|
||
completed += 1
|
||
if file_progress:
|
||
file_progress("file_chunk", fname,
|
||
chunks_done=completed, chunks_total=len(chunks))
|
||
all_ent = [ent for chunk_result in chunk_results for ent in chunk_result]
|
||
else:
|
||
all_ent = []
|
||
for k, c in enumerate(chunks):
|
||
try:
|
||
all_ent.extend(_call_llm(_build_llm_prompt(c, mapping.keys()), llm_client))
|
||
except Exception as e:
|
||
llm_failures += 1
|
||
log.warning("LLM NER chunk failed in %r: %s", fname, e)
|
||
if file_progress:
|
||
file_progress("file_chunk", fname,
|
||
chunks_done=k + 1, chunks_total=len(chunks))
|
||
|
||
# Дедуп + верификация + добавление в mapping
|
||
seen: set = set()
|
||
for ent in all_ent:
|
||
val = ent.get("value", "").strip()
|
||
if not val or val in mapping:
|
||
continue
|
||
nk = normalize_entity(val)
|
||
if nk in seen:
|
||
continue
|
||
seen.add(nk)
|
||
if not _verify_entity(val, full_text):
|
||
log.debug("LLM hallucination dropped: %r", val)
|
||
continue
|
||
raw_type = ent.get("type")
|
||
ent_type = raw_type.strip().lower().replace(" ", "_") if raw_type else ""
|
||
ent_type = ent_type or "other_id"
|
||
mapping[val] = _next_token(ent_type, counters)
|
||
|
||
if file_progress:
|
||
file_progress("file_done", fname, elapsed=round(time.time() - t0, 2))
|
||
|
||
return llm_failures
|