ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Настроить recurrent memory для long context

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

Реализовать механизм рекуррентной памяти для LLM, который позволяет модели сохранять важные факты между сессиями диалога без внешней RAG-системы. Память должна обновляться после каждого сеанса и встраиваться в контекст при старте нового. Ключевой результат Рабочий прототип чат-бота, способного запоминать информацию о пользователе на протяжении нескольких сессий (до 10+ фактов) и корректно отвечать на запросы, требующие этих знаний, с точностью не ниже 80%.

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

Если нет возможности запустить большую модель — симулируем:

1. Используем GPT-2 (small) — она достаточно легковесна.
2. Если нет GPU — работаем на CPU (инференс будет медленным, но для тестов достаточно).
3. Memory-модуль можно тестировать отдельно от LLM, подавая на вход эмбеддинги из sentence-transformers.

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

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

Этап 1: Исследование и проектирование архитектуры памяти (1–2 часа)

Действия

1. Изучите три подхода к долгосрочной памяти без RAG:
   • Memory Networks (Sukhbaatar et al., 2015) — явная память с вниманием.
   • Transformer-XL / Compressive Transformer — рекуррентные скрытые состояния.
   • Recurrent Memory — GRU/LSTM поверх эмбеддингов диалога.
2. Выберите один подход для реализации. Рекомендация гибрид — кодировать каждую сессию в вектор через GRU, хранить пул таких векторов, а при новом сеансе выбирать top-k по косинусной близости и вставлять в промпт.
3. Нарисуйте схему data flow:
   • Пользовательский ввод → LLM → ответ.
   • После сессии: конкатенация всего диалога → кодирование GRU → обновление памяти.
   • Перед новой сессией: поиск похожих векторов памяти → конвертация в текст → вставка в system prompt.

Ожидаемый результат этапа Документ (markdown) с выбором подхода, схемой и расчётом необходимого размера памяти (например, 512 dim, 10 слотов).

Этап 2: Реализация модуля рекуррентной памяти (2–3 часа)

Действия

1. Создайте класс RecurrentMemory:

   class RecurrentMemory(nn.Module):
       def __init__(self, input_dim=768, hidden_dim=512, num_slots=10):
           super().__init__()
           self.encoder = nn.GRU(input_dim, hidden_dim, batch_first=True)
           self.memory = nn.Parameter(torch.zeros(num_slots, hidden_dim))
           self.num_slots = num_slots
   
       def update(self, sequence_embeds):  # [batch, seq_len, input_dim]
           # seq_embeds == эмбеддинги предложений сессии
           _, h_n = self.encoder(sequence_embeds)
           # h_n: [1, batch, hidden_dim]
           new_mem = h_n.squeeze(0)  # [batch, hidden_dim]
           # Заменяем самый старый слот (простая очередь)
           self.memory = torch.roll(self.memory, shifts=-1, dims=0)
           self.memory[-1] = new_mem.detach()
           return self.memory
   
       def retrieve(self, query_embed, k=3):
           # query_embed: [1, hidden_dim] (кодировка начала новой сессии)
           scores = F.cosine_similarity(query_embed, self.memory, dim=-1)
           topk_indices = scores.topk(k).indices
           return self.memory[topk_indices]
   

2. Напишите юнит-тесты:
   • test_update_slots — проверка, что после 12 обновлений заполнились все слоты.
   • test_retrieve_returns_topk — проверка, что retrieve возвращает правильное количество.
   • test_memory_persistence — память сохраняется между вызовами.

Ожидаемый результат этапа Файл memory_module.py с классом RecurrentMemory и набором тестов (все проходят).

Этап 3: Интеграция памяти с LLM (1–2 часа)

Действия

1. Загрузите модель LLM (например, GPT2LMHeadModel from transformers).
2. Напишите функцию chat_with_memory(user_input, memory, model, tokenizer):
   • Если memory содержит слоты → декодируем top-3 воспоминания в текст (через предобученный декодер или просто храним сам текст).
   • Формируем system prompt: "You have the following memories: {memories}. User says: {input}".
   • Генерируем ответ.
   • После окончания сессии (по команде /end или timeout) вызываем memory.update(embeddings_of_session).
3. Реализуйте простой CLI-чат, который запускает сессии (каждая сессия — последовательность реплик).
4. Сохраняйте состояние памяти на диск (pickle) между запусками программы.

Ожидаемый результат этапа Скрипт chat_memory.py, который можно запустить, проверить сохранение факта через сессии.

Этап 4: Тестирование долгосрочной памяти (1–2 часа)

Действия

1. Составьте набор тестовых сценариев (не менее 5):
   • Сценарий A Сессия 1: "My name is Alice." / Сессия 2: "What is my name?" → должен ответить "Alice".
   • Сценарий B Сессия 1: "I work at Google." / Сессия 2: "Where do I work?" → "Google".
   • Сценарий C Пересечение: "I like pizza." / Сессия 3: "What is my favorite food?" → "pizza".
   • Сценарий D Множественные факты: 3 факта за 1 сессию, проверка в сессии 5.
   • Сценарий E Конфликт: в сессии 2 меняется факт (например, "I changed my name to Bob"), проверка, что новая информация перезаписывает старую.
2. Для каждого сценария записывайте:
   • prompt, ответ, ожидание.
   • Считайте accuracy = количество правильных ответов / общее число запросов.
3. Запустите тесты 5 раз для оценки стабильности.

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

Этап 5: Бенчмаркинг и сравнение с RAG-подходом (1 час)

Действия

1. Реализуйте простую RAG-альтернативу: те же факты хранятся в ChromaDB, поиск по косинусной близости query-эмбеддинга.
2. Сравните по трём метрикам:
   • Quality accuracy на том же тестовом наборе.
   • Latency среднее время ответа (сек).
   • Cost: дополнительное количество параметров / использование диска.
3. Сформируйте сравнительную таблицу.
4. Напишите вывод: в каких сценариях recurrent memory предпочтительнее RAG (например, низкий latency, приватность, offline).

Ожидаемый результат этапа Файл benchmark.pdf или benchmark.md с таблицей и выводами.

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

• Реализован модуль RecurrentMemory с методами update и retrieve.
• Модуль покрыт unit-тестами (минимум 3 теста, все проходят).
• Интегрированный чат-бот способен запомнить и извлечь до 10 независимых фактов из прошлых сессий.
• Accuracy на тестовом наборе сценариев ≥ 80%.
• Время инференса (с учётом памяти) не более 2x от baseline (без памяти).
• Память сохраняется между перезапусками программы (сериализация/десериализация).
• Создан отчёт с тестированием и сравнением с RAG.
• Код выложен в Git-репозиторий с README и инструкцией по запуску.

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

Основной артефакт Репозиторий со следующей структурой:

├── memory_module.py          # класс RecurrentMemory
├── chat_memory.py            # CLI-чат с интеграцией
├── tests/
│   ├── test_memory.py        # unit-тесты
│   └── test_scenarios.py     # end-to-end сценарии
├── report.md                 # отчёт тестирования и бенчмарк
├── requirements.txt
└── README.md

Содержимое отчёта accuracy по 5 сценариям, среднее время генерации, сравнение с RAG.

Опциональные дополнительные результаты

• Код для визуализации заполнения слотов памяти (matplotlib).
• Веса обученного memory encoder (если обучение было).
• Демо-видео работы чат-бота.

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

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

Примечание При первом выполнении возможны задержки на отладку и переключение между этапами — рекомендуем заложить +20% времени.

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

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

• Я реализовал класс RecurrentMemory с методами update и retrieve, и все юнит-тесты проходят.
• Я протестировал чат-бота на 5 сценариях и получил accuracy ≥ 80%.
• Я убедился, что память сохраняется между сессиями (факт из сессии 1 доступен в сессии 3).
• Я сравнил своё решение с RAG-подходом и зафиксировал метрики.
• Я написал отчёт и разместил код в Git-репозиторий с README.