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.

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

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