DE-1. DIY ассистент на LLM

Привет Хабр, let’s set the future.

Введение

Недавно у меня появилась идея фикс: ‘Хочу собственного AI ассистента’. Казалось бы, нет никаких проблем — рынок предлагает массу готовых решений. Но моя вечная паранойя про утечку данных и стремление сделать все самому взяли верх. Решил поэкспериментировать и собрать ассистента своими руками, да еще как-то с учетом будущих возможностей для гибкой настройки. Времени на оптимизацию производительности и эстетический вид кода у меня не было, ‘хочу здесь и сейчас’, поэтому let me introduce this shit.

Инструменты

Думаю, стоит сразу описать вкратце окружение:

• Для более эффективной работы в рамках linux окружения я использую WSL2 на Windows. На текущий момент используется дистрибутив Ubuntu-22.04.

• По поводу главного устройства, которое будет вычислять наши тензоры. GPU на 8gb (пример gtx 1080 и выше) должно хватить. На самом деле если очень не понятно где и как посмотреть требования выбранной вами LLM к памяти, то можно воспользоваться таким ПО как LM Studio.

• Чтобы все вычисления запустились на видеокарте, также стоит позаботиться о cuDNN драйверах. Тема установки стоит отдельной статьи, но благо такие уже есть: вариант 1 (все сам), вариант 2 (с помощью conda).

• Ollama — фреймворк для локального запуска крупных языковых моделей. Это то, что обязательно нужно для запуска ядра ассистента — LLM. Процесс установки фреймворка описан на официальном сайте.

Для реализации ассистента я выбрал три ключевые нейросети:

• STT — Whisper. Это модель для распознавания речи, разработанная OpenAI. Она способна обрабатывать аудиофайлы и переводить их в текст, поддерживает множество языков и может работать даже в условиях шума.

• LLM — Llama3. Это относительно новая LLM, по сравнению со своими предшественниками, она обладает улучшенной производительностью и более совершенными параметрами модели. Она способна отвечать на вопросы, предоставлять информацию и даже вести беседы на основе заданного контекста.

• TTS — Coqui AI. Система преобразования текста в речь, позволяет озвучивать текстовые ответы. Из всех open source решений предлагает достаточно естественное звучание и гибкость в настройках голоса и интонации на множестве языков.

Распознавание речи. Whisper

Приступим. Самый первый модуль необходим для преобразования голоса в текст, и для этой задачи отлично подошла модель Whisper. Она имеет несколько конфигураций: base, small, medium и large. Наилучшие результаты показывает модель base, которая обеспечивает оптимальный баланс между производительностью и качеством распознавания.
Функционал следующего кода очень прост. Внутри класса WhisperService происходит загрузка модели для преобразования аудио в текст с помощью библиотеки Whisper. Метод transcribe принимает путь к аудиофайлу в формате WAV и, используя модель, преобразует его в текст.

from abc import ABC, abstractmethod  class BaseService(ABC):     def __init__(self, model):         self.s2t = model      @abstractmethod     def transcribe(self, path_to_wav_file: str):         """         Abstract method to process audio files (in wav format) to text         """         pass   class WhisperService(BaseService):     _BASE_MODEL_TYPE = 'base'      def __init__(self, model_type: str = _BASE_MODEL_TYPE) -> None:         import whisper          model = whisper.load_model(model_type)         super().__init__(model)      def use_model(self, path_to_wav_file: str, language=None):         return self.s2t.transcribe(path_to_wav_file, language=language)      def transcribe(self, path_to_wav_file: str, language=None) -> str:         result = self.use_model(path_to_wav_file, language=language)         return result['text']

Обработка запросов. Llama3

Следующим важным звеном является модуль генерации текста. На текущий момент используется базовая LLM, в моей конфигурации _BASE_MODEL = llama3.1:latest. Код представленный ниже реализует модуль, который взаимодействует с языковой моделью с использованием библиотеки langchain_ollama. Основная цель модуля — отправка вопросов к модели и получение ответов. В методе ask_model, который отвечает за формирование запросов к модели, используется регулярное выражение для определения конца предложений. Метод получает вопрос, отправляет его в модель и обрабатывает потоковый ответ. Ответы накапливаются в буфере, и как только в буфере обнаруживается завершенное предложение, оно извлекается и возвращается. Таким образом, метод эффективно обрабатывает длинные ответы и позволяет как можно скорее передать созданное предложение в TTS модуль.

import re from langchain_ollama import ChatOllama  from config import LLM_MODEL   class LangChainService:     _BASE_MODEL = LLM_MODEL      def __init__(self, model_type: str = _BASE_MODEL):         self.model = ChatOllama(model=model_type)         self.context = ''      def ask_model(self, question: str):         buffer = ''         sentence_end_pattern = re.compile(r'[.!?]')          for chunk in self.model.stream(f'{self.context}\n{question}'):             buffer += str(chunk.content)             while True:                 match = sentence_end_pattern.search(buffer)                 if match:                     end_idx = match.end()                     sentence = buffer[:end_idx].strip()                     sentence = sentence[0 : len(sentence) - 1]                     yield sentence                     buffer = buffer[end_idx:].strip()                 else:                     break 

Синтез речи. Coqui AI

Ну и последний шаг, это преобразование ответа от бота в аудио формат. Этого можно достичь с помощью модуля для преобразования текста в речь, используя библиотеку XTTS. XTTSService инициализирует модель TTS, загружая её на доступное устройство, будь то GPU или CPU. Основная функция этого сервиса заключается в методе processing, который принимает текст и сохраняет его в виде аудиофайла формата WAV. Метод также позволяет указать язык и говорящего и скорость воспроизведения для более гибкой настройки.

from abc import ABC, abstractmethod  import torch  from config import TTS_XTTS_MODEL, TTS_XTTS_SPEAKER, TTS_XTTS_LANGUAGE   class BaseService(ABC):     def __init__(self, model):         self.t2s = model      @abstractmethod     def processing(self, text: str):         """         Abstract method to process text to audio files (in wav format)         """         pass   class XTTSService(BaseService):     _BASE_MODEL_TYPE = TTS_XTTS_MODEL     _BASE_MODEL_SPEAKER = TTS_XTTS_SPEAKER     _BASE_MODEL_LANGUAGE = TTS_XTTS_LANGUAGE      def __init__(self, model_type: str = _BASE_MODEL_TYPE) -> None:         from TTS.api import TTS          device = torch.device("cuda" if torch.cuda.is_available() else "cpu")         print(f'Apply {device} device for XTTS calculations')          model = TTS(model_type).to(device)          super().__init__(model)      def processing(         self,         path_to_output_wav: str,         text: str,         language: str = _BASE_MODEL_LANGUAGE,         speaker: str = _BASE_MODEL_SPEAKER,     ):         self.t2s.tts_to_file(text=text, file_path=path_to_output_wav, language=language, speaker=speaker, speed=2) 

Main.py скрипт. Telegram API

Чтобы быстро и без проблем собрать описанные выше модули и запустить ассистента, можно реализовать коммуникацию с ним через TelegramAPI. Плюсы: не нужно реализовывать клиента для записи и воспроизведения аудио. Минусы: не очень удобный UX, постоянно надо клацать кнопку записи в интерфейсе)

Telegram-бот разработан с использованием библиотеки python-telegram-bot.

Краткая логика работы:

1. Команда /start: Пользователь начинает взаимодействие с ботом, получая приветственное сообщение.

2. Обработка голосовых сообщений: Бот принимает голосовые сообщения от пользователей, проверяет их наличие и конвертирует в wav и сохраняет.

3. Распознавание речи: С помощью сервиса WhisperService аудиофайлы преобразуются в текст.

4. Генерация ответов: С помощью LangChainService текстовые команды обрабатываются, и генерируются текстовые ответы.

5. Преобразование текста в речь: Ответы преобразуются в голосовые сообщения с использованием XTTSService.

6. Отправка ответов: Генерированные голосовые сообщения отправляются обратно пользователю.

Ниже представлена простыня, которая реализует описанную выше логику:

from telegram import Update from telegram.ext import filters, Application, CommandHandler, CallbackContext, MessageHandler  from config import TELEGRAM_BOT_TOKEN  from src.generative_ai.services import LangChainService from src.speech2text.services import WhisperService from src.fs_manager.services import TelegramBotApiArtifactsIO from src.audio_formatter.services import PydubService from src.text2speech.services import XTTSService from src.telegram_api.services import user_verification from src.shared.hash import md5_hash   speech_to_text = WhisperService() text_to_speech = XTTSService() file_system = TelegramBotApiArtifactsIO() formatter = PydubService() langchain = LangChainService()   async def verify_user(update: Update) -> None:     user_id: str = str(update.effective_user.id)  # type: ignore     user_verification(user_id)   async def start(update: Update, _: CallbackContext) -> None:     await verify_user(update)     await update.message.reply_text('Hello! I am your personal assistant. Let is start)')  # type: ignore   async def handle_audio(update: Update, context: CallbackContext) -> None:     await verify_user(update)      artifact_paths = []      user_id: str = str(update.effective_user.id)  # type: ignore     chat_id = update.message.chat_id  # type: ignore     voice_message = update.message.voice  # type: ignore      if not voice_message:         await update.message.reply_text('Please, send me audio file.')  # type: ignore         return      input_file_path = await file_system.write_user_audio_file(user_id, voice_message)     artifact_paths.append(input_file_path)     output_file_path = formatter.processing(input_file_path, '.wav')  # type: ignore     artifact_paths.append(output_file_path)     text_message = speech_to_text.transcribe(output_file_path)      for text_sentence in langchain.ask_model(text_message):         sentence_hash = md5_hash(text_sentence)         wav_ai_answer_filepath = file_system.make_user_artifact_file_path(             user_id=user_id, filename=f'{sentence_hash}.wav'         )         artifact_paths.append(wav_ai_answer_filepath)         text_to_speech.processing(wav_ai_answer_filepath, text_sentence)         ogg_ai_answer_filepath = formatter.processing(wav_ai_answer_filepath, '.ogg')         artifact_paths.append(ogg_ai_answer_filepath)         await send_voice_message(context=context, chat_id=chat_id, file_path=ogg_ai_answer_filepath)      file_system.delete_artifacts(user_id=user_id, filename_array=artifact_paths)   async def send_voice_message(context: CallbackContext, chat_id, file_path: str):     with open(file_path, 'rb') as voice_file:         await context.bot.send_voice(chat_id=chat_id, voice=voice_file)   def main() -> None:     application = Application.builder().token(TELEGRAM_BOT_TOKEN).build()      application.add_handler(CommandHandler('start', start))     application.add_handler(MessageHandler(filters.VOICE & ~filters.COMMAND, handle_audio))      application.run_polling()   if __name__ == '__main__':     main()

Браузерный клиент. WebSockets

После работы с Telegram-ботом я пришёл к выводу, что его функционал не совсем удобен при реализации полноценного голосового ассистента. Бот хотя и предоставляет базовые возможности взаимодействия, ограничивает меня в плане пользовательского опыта. Поэтому я стал думать, как можно менее болезненно реализовать клиентское приложение.
Самым очевидным решением для меня оказался браузерный клиент на основе WebSocket. Плюсы: подключение устройств записи и воспроизведения звука через браузер, возможность реализации клиента на любом устройстве.

Вот такой клиент получился на скорую руку. Здесь все просто записанные фреймы на постоянной основе шлются на бек, в то время как аудио ответы собираются в очередь и синхронно воспроизводятся с помощью функции playNextAudio. Ниже представлен код клиента:

<!DOCTYPE html> <html lang="en"> <head>     <meta charset="UTF-8" />     <meta name="viewport" content="width=device-width, initial-scale=1.0" />     <title>Chekov</title> </head> <body>     <button id="startBtn">Start Recording</button>     <button id="stopBtn">Stop Recording</button>     <button id="enableAudioBtn">Enable Audio Playback</button>     <script>         const TIME_SLICE = 100;         const WS_HOST = "localhost";         const WS_PORT = 8765;          const WS_URL = `ws://${WS_HOST}:${WS_PORT}`;         const ws = new WebSocket(WS_URL);          ws.onopen = () => console.log("WebSocket connection established");         ws.onclose = () => console.log("WebSocket connection closed");         ws.onerror = (e) => console.error("WebSocket error:", e);         ws.onmessage = (event) => collectVoiceAnswers(event);          let mediaRecorder;         let audioEnabled = false;         const audioQueue = [];         let isPlaying = false;          async function startRecord() {             const userMediaSettings = { audio: true };             const stream = await navigator.mediaDevices.getUserMedia(userMediaSettings);             mediaRecorder = new MediaRecorder(stream);             mediaRecorder.ondataavailable = streamData;             mediaRecorder.start(TIME_SLICE);         }          function streamData(event) {             if (event.data.size > 0 && ws.readyState === WebSocket.OPEN) {                 wsSend(event.data);             }         }          function wsSend(data) {             ws.send(data);         }          function stopRecord() {             if (mediaRecorder) {                 mediaRecorder.stop();             }         }          function collectVoiceAnswers(event) {             if (!audioEnabled) return;              const { data, type } = JSON.parse(event.data);             const audioData = atob(data);             const byteArray = new Uint8Array(audioData.length);              for (let i = 0; i < audioData.length; i++) {                 byteArray[i] = audioData.charCodeAt(i);             }              const audioBlob = new Blob([byteArray], { type: "audio/wav" });             const audioUrl = URL.createObjectURL(audioBlob);             const audio = new Audio(audioUrl);             audioQueue.push({ audio, type });             if (!isPlaying) {                 playNextAudio();             }         }          async function playNextAudio() {             if (audioQueue.length === 0) {                 isPlaying = false;                 return;             }              isPlaying = true;             const { audio, type } = audioQueue.shift();             try {                 await new Promise((resolve, reject) => {                     audio.onended = resolve;                     audio.onerror = reject;                     audio.play().catch(reject);                 });                 playNextAudio();             } catch (e) {                 console.error("Error playing audio:", e);                 isPlaying = false;             }         }          document.getElementById("startBtn").addEventListener("click", startRecord);         document.getElementById("stopBtn").addEventListener("click", stopRecord);     </script> </body> </html> 

Реализацию серверной части, которая обрабатывает WebSocket-соединения и взаимодействует с остальной частью ассистента, вы можете получить по указанной ссылке. Также по этой ссылке в репозитории можно найти quick start guide.

Заключение

Вот и все. Для дальнейшего улучшения ассистента, включая добавление новых функций(именно функций для ассистирования, чтобы бот начал оправдывать свое название), таких как сохранение заметок, поиск информации и другие полезные фичи, стоит рассмотреть возможность файн-тюнинга LLM, чтобы выдавать унифицированные ответы в формате {command, message}. Также полезным будет реализация постпроцессинга для обработки команд с использованием классических алгоритмов на основе вывода LLM.

А на этом все. Спасибо, что дочитали до конца!

Тут оставлю ссылку на весь код ассистента