Longformer chunk pruner для Q-RAG

Модель от ArtTolston на базе allenai/longformer-base-4096: Longformer с обученной линейной головкой keep/drop. Она получает вопрос и цепочку чанков из Q-RAG, оценивает каждый чанк вероятностью p_keep и по умолчанию оставляет те, где p_keep >= 0.5. Это кастомная архитектура прунинга, а не стандартная sequence-classification из Transformers.

📌 Ключевые особенности

• ✅ Оценивает каждый retrieved-чанк: keep_probabilities и kept_indices в одном вызове.
• ✂ Не выполняет retrieval: работает поверх готовой цепочки Q-RAG и требует исходный порядок и формат чанков.
• 👀 Вход: вопрос до 128 токенов, каждый чанк до 512, вся сборка до 4096 токенов; глобальное внимание на [CLS] и токенах вопроса.
• 📌 Чанк, который не влезает в лимит, получает None и исключается.
• ⭐ Обучена на HotpotQA и MuSiQue; есть готовый loader load_pruner и функция prune.

💻 Пример

import sys
from huggingface_hub import snapshot_download
model_dir = snapshot_download("ArtTolston/longformer-base-pruner-hotpotqa_musique")
sys.path.insert(0, model_dir)
from inference import load_pruner, prune
model, tokenizer = load_pruner(model_dir, device="cpu")
result = prune(model, tokenizer, question, chunks, threshold=0.5)
print(result["keep_probabilities"], result["kept_indices"])

Подробнее

На отобранном чекпоинте: supporting-fact F1 0.9546 против 0.5149 без прунинга; precision 0.9666, recall 0.9492.
Покрытие всех gold-support чанков — 0.8867 против 0.9800 без прунинга.
Это метрики на 300-вопросной выборке для выбора чекпоинта, не полноценный held-out бенчмарк.
148,660,225 параметров; запуск возможен на CPU или CUDA, VRAM не указан.