NER для X5 Retail Group

📋 О проекте

Сервис выделения сущностей из поискового запроса клиента в мобильном приложении торговой сети «Пятерочка»

Задача: извлечение из текстовых описаний товаров следующих сущностей:

• VOLUME - объем/вес товара (например: "500ml", "1кг")
• TYPE - тип продукта (например: "молоко", "сыр")
• PERCENT - процентное содержание (например: "3.2%", "жирность 15%")
• BRAND - бренд/производитель (например: "Простоквашино", "Parmalat")

Проект включает несколько подходов к решению задачи NER с различными архитектурами и оптимизациями, а также готовые API для развертывания на сервере.

---

🏗️ Архитектура моделей

Детальное описание архитектур

1️⃣ Базовая модель (BERT)

Tokens → BERT Encoder → Dropout → Linear(hidden_size → num_labels) → Softmax

Особенности:

• AutoModelForTokenClassification из Transformers
• Независимое предсказание для каждого токена (не учитывает зависимости между метками)
• Baseline для сравнения

2️⃣ BERT + Weighted Loss

Tokens → BERT Encoder → Dropout → Linear → Softmax + WeightedCrossEntropyLoss

Особенности:

• Идентична базовой модели, но с class weights в функции потерь
• Решает проблему дисбаланса классов
• Weights = inverse frequency для каждого класса

3️⃣ LLM + BiLSTM + CRF

Text → LLM Tokenizer → Frozen LLM
     → Word Embeddings → BiLSTM(2 layers, hidden=256) 
     → Dropout(0.5) → Linear → CRF

Особенности:

• Использует контекстные эмбеддинги из замороженной LLM (feature extraction)
• BiLSTM моделирует последовательность в обоих направлениях
• CRF обеспечивает валидность последовательности меток
• Устойчивость к опечаткам и OOV (out-of-vocabulary) словам

4️⃣ BERT + CRF

Tokens → BERT Encoder → Dropout(0.1) 
       → Linear(hidden_size → num_labels) 
       → CRF(emissions + transitions) → Viterbi Decoding

Особенности:

• Кастомная модель BertCrfForTokenClassification с CRF слоем из библиотеки torchcrf
• CRF моделирует матрицу переходов между метками для структурного предсказания
• Viterbi алгоритм находит оптимальную последовательность меток
• Loss = -log likelihood всей последовательности (CRF обучается end-to-end)

5️⃣ BERT + CRF + Additional Features

Tokens → BERT Encoder → Dropout(0.1)
       → [BERT_hidden ⊕ Lang_Embedding(3→8) ⊕ Length(1)] 
       → Linear(hidden_size+9 → num_labels) → CRF → Viterbi

Особенности:

• Идентична src_crf, но с расширенным входом в Linear слой
• Lang Type Embedding(3→8): кодирование типа символов (0=латиница, 1=кириллица, 2=цифры)
• Word Length [0,1]: нормализованная длина токена как дополнительный признак
• Итоговая размерность: hidden_size + 8 + 1 = hidden_size + 9

---

📁 Структура проекта

Xak_5_NER/
├── src_base/              # Базовая BERT модель
│   ├── train_base.py
│   ├── inference_base.py
│   └── NER_model_train_base_pipline.ipynb
├── src_base_v1/           # Базовая BERT модель (v1)
│   ├── train_v1.py
│   ├── inference_base.py
│   └── NER_model_train.ipynb
├── src_bilstm/            # BERT + BiLSTM
│   ├── train_bilstm.py
│   ├── inference_bilstm.py
│   └── NER_model_train_v2_BiLSTM.ipynb
├── src_crf/               # BERT + CRF
│   ├── train crf.py
│   ├── inference_crf.py
│   └── NER_model_train_v1_crf.ipynb
├── src_crf_v1/            # BERT + CRF (v1)
│   ├── train.py
│   ├── inference.py
│   └── NER_model_train_v3_crf_add.ipynb
├── data/                  # Датасеты
│   ├── train_balanced.csv
│   ├── train_v1.csv
│   ├── train_extended.csv
│   └── dataset_generation.ipynb
│   └── balance_dataset.py
├── API/                   # API для базовой модели
│   └── SRC/
│       ├── app.py
│       ├── inference.py
│       ├── Dockerfile
│       └── model_finetuned/    # ⚠️ Модель загружается отдельно
└── API_crf/               # API для CRF модели
    └── SRC/
        ├── app.py
        ├── inference.py
        ├── Dockerfile
        └── model_finetuned/    # ⚠️ Модель загружается отдельно

---

🚀 Развертывание API

Описание

Папки API/ и API_crf/ содержат готовые к развертыванию FastAPI приложения с оптимизацией для production:

• Многопоточная обработка с uvicorn workers
• Высокая производительность (до 10,000+ RPS на 32 vCPU)
• Docker-контейнеризация с оптимизированными ресурсами
• Детальная документация внутри каждой папки

⚠️ Требования перед запуском

Важно! Перед развертыванием необходимо загрузить обученную модель с Hugging Face Hub:

1. Скачайте модель с Hugging Face:
   • Для базовой модели → P0ve1/NER_X5_bert
   • Для CRF модели → P0ve1/NER_X5_bert-crf
2. Распакуйте файлы модели в соответствующую папку:
   • Для базовой модели → API/SRC/model_finetuned/
   • Для CRF модели → API_crf/SRC/model_finetuned/

Требуемые файлы модели:

model_finetuned/
├── config.json
├── pytorch_model.bin (или model.safetensors)
├── tokenizer_config.json
├── tokenizer.json
├── vocab.txt
└── special_tokens_map.json

Запуск API

# Для базовой модели
cd API
docker-compose up -d --build

# Для CRF модели
cd API_crf
docker-compose up -d --build

Примеры использования

# Health check
curl http://localhost:80/health

# Предсказание (базовая модель)
curl -X POST "http://localhost:80/api/predict" \
  -H "Content-Type: application/json" \
  -d '{"input": "молоко сгущенное 500ml"}'

# Production endpoint (пример)
curl -X POST "http://51.250.30.54:8000/api/predict" \
  -H "Content-Type: application/json" \
  -d '{"input": "молоко сгущенное 500ml"}'

Ответ:

[
  {"start_index": 0, "end_index": 6, "entity": "B-TYPE"},
  {"start_index": 7, "end_index": 17, "entity": "B-BRAND"},
  {"start_index": 18, "end_index": 23, "entity": "B-VOLUME"}
]

📖 Документация API

Подробные инструкции по настройке, оптимизации и мониторингу см. в файлах:

• API/README.md - полная документация для базовой модели
• API/instruction.md - пошаговая инструкция по развертыванию
• API_crf/README.md - документация для CRF модели
• API_crf/instruction.md - инструкция для CRF версии

Документация включает:

• Расчет оптимальных параметров workers/concurrency для вашей конфигурации
• Готовые конфигурации для серверов от 4 до 64 vCPU
• Мониторинг и отладка производительности
• Рекомендации по тонкой настройке

---

💾 Данные

Датасеты

Папка data/ содержит:

• train_v1.csv - исходный датасет
• train_balanced.csv - сбалансированный датасет для обучения
• train_extended (3).csv - расширенный датасет
• balance_dataset.py - скрипт для балансировки данных

Формат данных

Датасет представлен в CSV формате со следующими колонками:

• sample - текст описания товара (токенизированный по пробелам)
• annotation - список аннотаций в формате [(token, label), ...]

Пример:

sample,annotation
"молоко сгущенное 500ml","[('молоко', 'B-TYPE'), ('сгущенное', 'I-TYPE'), ('500ml', 'B-VOLUME')]"

Метки (Labels)

Используется схема BIO (Begin, Inside, Outside):

• B-VOLUME / I-VOLUME - объем/вес
• B-TYPE / I-TYPE - тип продукта
• B-PERCENT / I-PERCENT - процентное содержание
• B-BRAND / I-BRAND - бренд
• O - не является сущностью

📊 Результаты моделей (Legacy)

Результаты старых экспериментов на выборке 5000 примеров:

Тестирование разных предобученных моделей (LR = 2e-5, num_epochs = 3, data=train_v1.csv)

Custom Optimizer (LR = 2e-5, num_epochs = 3, data=train_v1.csv)

---

🔧 Технологии

• PyTorch - фреймворк для обучения моделей
• Transformers (Hugging Face) - предобученные языковые модели
• torchcrf - реализация CRF слоя
• FastAPI - веб-фреймворк для API
• Uvicorn - ASGI сервер с поддержкой многопоточности
• Docker - контейнеризация приложений
• Accelerate - распределенное обучение

---

👨‍💻 Разработка

Структура кода

Каждая папка src_* содержит:

• train*.py - скрипт для обучения модели
• inference*.py - скрипт для предсказаний
• *.ipynb - Jupyter notebook с полным pipeline

Основные классы и функции

• DataPreprocessor - загрузка и предобработка данных
• align_labels_with_tokens - выравнивание меток с токенами
• compute_metrics - вычисление F1-score и accuracy
• BertCrfForTokenClassification - кастомная модель BERT+CRF
• LLMEmbeddingExtractor - извлечение эмбеддингов для BiLSTM