ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Fine-tune embedding для юридического домена

1. Цель задачи

Разработать и обучить векторную модель (embedding model), специализированную для поиска и сравнения юридических договоров. Используя contrastive loss на корпусе из 5000 договоров (с парами «похожих» и «непохожих» документов), улучшить качество поиска по сравнению с универсальным embedding-моделями. Ключевой результат повышение метрики Recall@10 на 20% (относительных) относительно базовой модели (например, all-MiniLM-L6-v2).

2. Исходные данные

Если нет реального инструмента – симулируем:

1. Возьмите 5000 любых длинных текстов (например, новости из AG_NEWS или случайные статьи из Wikipedia).
2. Разбейте каждый текст на 2–3 части (500–1000 токенов), считая части одного документа «похожими» (positive pair), а части разных документов – «непохожими» (negative pair).
3. Для синтеза договоров можно использовать генеративный LLM (GPT-4/Claude) с промптом «Напиши краткий договор аренды (поставки, подряда)» – 1 запрос = 1 договор ~500 токенов.

3. Технологический стек

4. Этапы выполнения

Этап 1: Подготовка данных (2 часа)

Действия

1. Собрать 5000 договоров (текстов). Если используете CUAD:

   from datasets import load_dataset
   dataset = load_dataset("cuad", split="train")
   contracts = dataset["context"]  # список текстов
   

   Для синтетического датасета – выполнить генерацию через LLM (см. п.2).

2. Разбить каждый договор на чанки по 256 слов (перекрытие 32 слова) – чанк станет единицей поиска.

3. Создать пары для обучения:

   • Positive pairs два чанка из одного договора (выбираем случайные соседние чанки).
   • Negative pairs чанк из договора A + чанк из договора B (разные договоры).
   • Итоговый набор: 20k positive + 20k negative (или по 50k – зависит от числа чанков).

4. Разделить на train/val/test (80/10/10) на уровне договоров (чтобы один договор не попал и в train, и в test).

Ожидаемый результат этапа

• Директория data/ с файлами train_pairs.csv, val_pairs.csv, test_pairs.csv (столбцы: anchor, positive, negative).
• Чек-суммы или хеши для верификации.

Этап 2: Baseline оценка (1 час)

Действия

1. Загрузить модель all-MiniLM-L6-v2:

   from sentence_transformers import SentenceTransformer
   model = SentenceTransformer('all-MiniLM-L6-v2')
   

2. Закодировать все чанки из test-набора (около 1500 чанков) эмбеддингами (dim=384).

3. Для каждого query (случайный чанк из test-набора) вычислить Recall@10:

   from sklearn.metrics.pairwise import cosine_similarity
   import numpy as np
   
   queries = ... # эмбеддинги 200 случайных чанков
   corpus = ... # все эмбеддинги test-набора
   # для каждого query найти 10 ближайших в корпусе
   # Recall@10 = доля запросов, среди 10 ближайших есть хотя бы один чанк из того же договора
   

   Зафиксировать baseline Recall@10.

Ожидаемый результат этапа

• Baseline Recall@10 (например, 0.45).
• Скрипт evaluate_baseline.py.

Этап 3: Fine-tune с contrastive loss (3–4 часа)

Действия

1. Настроить DataLoader на train_pairs.csv. Использовать InputExample из sentence-transformers:

   from sentence_transformers import InputExample
   
   def create_examples(row):
       return InputExample(texts=[row['anchor'], row['positive'], row['negative']], label=0)
   

2. Определить loss:

   from sentence_transformers.losses import MultipleNegativesRankingLoss
   train_loss = MultipleNegativesRankingLoss(model)
   

3. Обучить модель:

   from sentence_transformers import SentenceTransformer, SentencesDataset, losses
   model = SentenceTransformer('all-MiniLM-L6-v2')
   train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)
   model.fit(train_objectives=[(train_dataloader, train_loss)],
             epochs=5,
             warmup_steps=100,
             evaluator=... # опционально
             output_path='output/fine-tuned-legal')
   

4. Гиперпараметры:

   • batch_size: 16 (если позволяет GPU) или 8
   • learning_rate: 2e-5 (AdamW)
   • epochs: 5 (с early stopping по val loss)
   • max_seq_length: 256 токенов

5. Логировать loss каждые 10 шагов (wandb/tensorboard).

Ожидаемый результат этапа

• Сохранённая fine-tuned модель в output/fine-tuned-legal/.
• История обучения (loss, val loss) в виде графика.

Этап 4: Оценка и сравнение (1 час)

Действия

1. Повторить Этап 2 для fine-tuned модели (загрузить SentenceTransformer('output/fine-tuned-legal')).

2. Вычислить Recall@10 на том же test-наборе, что и baseline.

3. Рассчитать относительное улучшение: (new - baseline) / baseline * 100%.

4. Дополнительно: Precision@10, поиск ближайших соседей (визуализация через t-SNE).

Ожидаемый результат этапа

• Recall@10 improved >= +20% (например, с 0.45 до 0.54).
• Краткий отчёт evaluation_report.md.

Этап 5: Документирование и упаковка (1 час)

Действия

1. Оформить README: описание задачи, инструкция по запуску, примеры использования.

2. Создать requirements.txt.

3. Зафиксировать результаты в results.csv (модель, Recall@10, Precision@10, время обучения).

Ожидаемый результат этапа

• Репозиторий с воспроизводимыми скриптами.
• Инференс-скрипт embed_and_search.py для демонстрации.

5. Критерии приемки (Definition of Done)

• Подготовлен корпус из 5000 договоров (или синтетический эквивалент) с разбивкой на чанки.
• Созданы train/val/test пары для contrastive loss (минимум 10k пар).
• Baseline‑модель (all-MiniLM-L6-v2) оценена — получен Recall@10.
• Fine-tune завершён без ошибок (сходимость loss, модель сохранена).
• Recall@10 fine-tuned модели улучшился на ≥20% относительно baseline.
• Все скрипты запускаются одной командой (python train.py, python evaluate.py) на чистом окружении.
• Сгенерирован отчёт с графиками и метриками.
• Код опубликован в репозитории (Git) с README и requirements.txt.

6. Ожидаемый результат

• Основной артефакт директория fine-tuned-legal/ (веса модели, конфиг, tokenizer).
• Содержимое файлы model.safetensors, config.json, tokenizer.json, train_log.txt, evaluation_results.json.
• Отчёт evaluation_report.md (таблицы сравнения с baseline, графики Recall@10, выводы).
• Дополнительно скрипт search_demo.py, который по заданному тексту договора находит топ-10 похожих чанков из корпуса.

7. Возможные сложности и их решение

8. Бюджет времени (оценка)

Примечание для первого выполнения с учетом настройки окружения и отладки заложите +3 часа (до 12 часов). Обучение на 5000 чанках занимает ~30 минут на GPU T4.

9. Связанные вопросы из базы знаний

10. Чек-лист самопроверки

• Я проверил, что train и test наборы не пересекаются по договорам.
• Я обучил модель на случайных данных (миниатюра) перед полным запуском, чтобы убедиться, что код работает.
• Я сравнил Recall@10 на идентичном test-наборе для baseline и fine-tune модели.
• Я сохранил модель и конфиги в отдельную папку с версией.
• Я задокументировал шаги воспроизведения в README, включая команды и ожидаемое время выполнения.