#!/usr/bin/env bash
# Загрузка моделей распознавания для второй ступени.
#
# Используются две модели под разные задачи:
#   - 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")/.."

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
}

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

# LLAIM Legal NER (LLAIMlegal/ru-legal-ner, MIT) — юридические реквизиты и
# документы: ИНН, ОГРН, СНИЛС, паспорт, телефон, email, банковский счёт, дата.
# Готового ONNX нет — скачиваем веса и конвертируем скриптом (нужны torch,
# transformers, optimum, onnx).
fetch "https://huggingface.co/LLAIMlegal/ru-legal-ner/resolve/main" \
  models/ru-legal-ner model.safetensors \
  config.json tokenizer.json tokenizer_config.json

if command -v python3 >/dev/null 2>&1; then
  python3 tools/convert_ru_legal_ner.py || echo "Не удалось сконвертировать ru-legal-ner в ONNX"
else
  echo "python3 не найден: ru-legal-ner не сконвертирован в ONNX"
fi

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"
echo "  pdguard.ner.legal-engine=ru-legal-ner, pdguard.ner.legal-model=models/ru-legal-ner"