ТЕХНИЧЕСКОЕ ЗАДАНИЕ: RAG с cross-encoder reranking

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

Реализовать пайплайн двухэтапного retrieval для RAG: сначала ANN (приблизительный поиск ближайших соседей) возвращает top-100 кандидатов, затем cross-encoder переранжирует их в top-10. Измерить качество ранжирования метрикой NDCG@10 и добиться её улучшения на 30% относительно baseline (только ANN, без реранжирования). Задача развивает навыки интеграции эффективного поиска с точным реранжированием и оценки ранжирующих систем.

Ключевой результат Пайплайн ANN → cross-encoder → top-10 с фиксированным датасетом, где NDCG@10 улучшен на ≥30% по сравнению с baseline.

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

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

1. Используем datasets.load_dataset("ms_marco", "passage_ranking") — это dev queries (6980 запросов) и пассажи с ID.
2. Если qrels не хватает, возьмём первые 1000 запросов и вручную не надо — qrels уже есть в датасете.
3. Если нет GPU для cross-encoder, используем меньшую модель: cross-encoder/stsb-distilroberta-base (но лучше ms-marco) или batch inference на CPU (медленно, но допустимо для небольшой выборки). Для оценки NDCG достаточно прогнать 200–500 запросов.

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

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

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

Действия

1. Загрузить датасет MS MARCO passage ranking:

   from datasets import load_dataset
   dataset = load_dataset("ms_marco", "passage_ranking", split="dev")
   queries = dataset.select(range(1000))  # уменьшить для скорости
   

2. Извлечь ID запросов, текст запросов, ID релевантных документов (qrels).
3. Загрузить корпус пассажей (corpus) — отдельный датасет ms_marco.passages.
4. Построить эмбеддинги пассажей с bi-encoder:

   from sentence_transformers import SentenceTransformer
   bi_encoder = SentenceTransformer('all-MiniLM-L6-v2')
   corpus_texts = [passage['text'] for passage in corpus]
   corpus_embeddings = bi_encoder.encode(corpus_texts, convert_to_numpy=True, show_progress_bar=True)
   

5. Построить FAISS индекс (FlatIP — внутреннее произведение, т.е. cosine после нормализации):

   import faiss
   index = faiss.IndexFlatIP(corpus_embeddings.shape[1])
   faiss.normalize_L2(corpus_embeddings)
   index.add(corpus_embeddings)
   

6. Baseline retrieval для каждого запроса получить top-10 по ANN. Для этого нормализовать эмбеддинг запроса (L2) и выполнить index.search.
7. Вычислить NDCG@10 baseline:

   from sklearn.metrics import ndcg_score
   # подготовка релевантностей: 1 для релевантных, 0 для нерелевантных
   results_baseline = []
   for qid, q_emb in zip(query_ids, query_embeddings):
       distances, indices = index.search(q_emb[np.newaxis,:], 10)
       doc_ids = [corpus_ids[i] for i in indices[0]]
       # получить метки релевантности
       y_true = [1 if did in qrels[qid] else 0 for did in doc_ids]
       y_score = [1 - distances[0][i] for i in range(10)]  # чем больше distance, тем ближе; инвертируем
       ndcg_baseline.append(ndcg_score([y_true], [y_score], k=10))
   print("Baseline NDCG@10:", np.mean(ndcg_baseline))
   

Ожидаемый результат этапа Файл baseline_metrics.json со средним NDCG@10 и Recall@10.

Этап 2: Реализация re-ranker (cross-encoder) (1.5 часа)

Действия

1. Загрузить cross-encoder модель:

   from sentence_transformers import CrossEncoder
   model_ce = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', max_length=512)
   

2. Для каждого запроса получить top-100 кандидатов через ANN (из FAISS).
   • Записать пары (запрос, пассаж) для всех кандидатов.
   • Внимание: не перепутать нормализацию; ANN должен вернуть 100 кандидатов.
3. Прогнать cross-encoder на парах (query, passage) батчами (batch_size=32–64) для ускорения.
4. Отсортировать кандидаты по убыванию скоров cross-encoder, взять top-10.
5. Вычислить NDCG@10 для re-ranked результатов.

Ожидаемый результат этапа Файл reranked_metrics.json с метриками.

Этап 3: Сравнение и анализ (1 час)

Действия

1. Сравнить средний NDCG@10 baseline и re-ranked.
2. Вычислить относительное улучшение: (ndcg_rerank - ndcg_baseline) / ndcg_baseline * 100.
3. Построить график распределения NDCG@10 по запросам (гистограмма или boxplot).
4. Проанализировать случаи, где реранжирование ухудшило результат (если есть). Возможные причины: шум в qrels, неудачная модель cross-encoder.
5. Зафиксировать дополнительно Recall@10, MRR@10 (для полноты).

Ожидаемый результат этапа Отчёт comparison_report.md с таблицами метрик и графиками.

Этап 4: Оптимизация (опционально, 1 час – при желании улучшить результат)

Действия

1. Попробовать другой cross-encoder: cross-encoder/ms-marco-TinyBERT-L-2 (быстрее) или BAAI/bge-reranker-v2-m3 (если доступен).
2. Увеличить число ANN кандидатов (top-200 вместо 100) — может дать прирост, но замедлит.
3. Применить hybrid: weighted sum ANN score + cross-encoder score.
4. Повторить этап 2 и 3, записать улучшения.

Ожидаемый результат этапа Файл optimized_metrics.json с лучшим результатом.

Этап 5: Фиксация и рефакторинг (0.5 часа)

Действия

1. Структурировать код в модули: embed.py, search.py, rerank.py, evaluate.py.
2. Написать main.py, который запускает весь пайплайн с параметрами (bi-encoder, cross-encoder, top_k).
3. Добавить README.md с инструкцией по запуску и результатами.
4. Закоммитить в Git.

Ожидаемый результат этапа Чистый репозиторий с воспроизводимым пайплайном.

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

• Датасет загружен и обработан (корпус, запросы, qrels).
• Реализован baseline retrieval (ANN → top-10) с измеренным NDCG@10.
• Реализован re-ranking пайплайн (ANN top-100 → cross-encoder → top-10).
• NDCG@10 после реранжирования улучшен на ≥30% относительно baseline.
• Все метрики зафиксированы в JSON-файлах.
• Код воспроизводим: однократный запуск python main.py выдаёт метрики.
• README.md содержит описание, команды запуска и итоговые метрики.
• Отчёт с визуализацией (гистограмма распределения NDCG) присутствует.
• Код залит в репозиторий с хотя бы одним коммитом.

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

• Репозиторий (папка проекта) с файлами:
  • main.py — точка входа.
  • embed.py — функции построения эмбеддингов и FAISS индекса.
  • search.py — ANN retrieval.
  • rerank.py — cross-encoder re-ranking.
  • evaluate.py — вычисление NDCG, Recall, MRR.
  • requirements.txt — зависимости.
  • README.md — описание, результаты, команды запуска.
  • baseline_metrics.json, reranked_metrics.json — JSON с метриками.
  • comparison_report.md — отчёт с таблицами и графиками.
• Метрики:
  • Baseline NDCG@10 = X (конкретное число, например 0.35)
  • Reranked NDCG@10 = Y (например 0.46)
  • Улучшение = (Y - X)/X * 100% ≥ 30%
• Опционально файл с визуализацией (PNG/PDF).

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

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

Примечание Для первого раза ожидаемый диапазон 4–6 часов в зависимости от опыта и производительности машины.

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