#!/usr/bin/env bash
# Обучение модели для второй ступени распознавания имён.
#
# Модель в репозиторий не кладётся: она весит мегабайты и собирается из открытых
# корпусов за несколько минут. Без модели сервис работает на одних правилах.
#
#   ./tools/train-ner.sh              # быстрый вариант, только factRuEval
#   ./tools/train-ner.sh full         # плюс префикс Nerus, качество заметно выше
set -euo pipefail
cd "$(dirname "$0")/.."

MODE="${1:-quick}"
WORK="$(mktemp -d)"
trap 'rm -rf "$WORK"' EXIT

echo "1. factRuEval-2016 — ручная разметка, 1965 предложений"
curl -sSL -o "$WORK/fre.tar.gz" \
  https://codeload.github.com/dialogue-evaluation/factRuEval-2016/tar.gz/refs/heads/master
tar xzf "$WORK/fre.tar.gz" -C "$WORK"
python3 tools/factrueval-to-opennlp.py "$WORK/factRuEval-2016-master" "$WORK/train.txt"

if [ "$MODE" = "full" ]; then
  echo "2. Nerus — автоматическая разметка, берём префикс потоком (400 тыс. предложений)"
  curl -sS -r 0-400000000 \
    https://storage.yandexcloud.net/natasha-nerus/data/nerus_lenta.conllu.gz \
    | gunzip 2>/dev/null \
    | python3 tools/conllu-to-opennlp.py 400000 >> "$WORK/train.txt"
fi

echo "3. Обучение, несколько минут"
mvn -q -B test-compile
mvn -q -B dependency:build-classpath -Dmdep.outputFile="$WORK/cp.txt"
mkdir -p models
java -Xmx6g -Dner.iterations=100 -Dner.cutoff=5 \
  -cp "target/test-classes:target/classes:$(cat "$WORK/cp.txt")" \
  ru.pdguard.tools.NerTrainer "$WORK/train.txt" models/ru-ner-person.bin

echo
echo "Готово. Включить: pdguard.ner.model=models/ru-ner-person.bin"