DE-1. DIY ассистент на LLM26.10.2024 20:15

Привет Хабр, 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. Ниже представлен код клиента:

    
    
    


    Start Recording
    Stop Recording
    Enable Audio Playback
    

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

Заключение

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

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

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

© Habrahabr.ru