Стою я значит утром (около 2 часов дня) возле кофеварки и листаю ленту хабра, а там CodeLama вышла. Copilot для бедных это или панацея в мире локальных текстовых моделей? Попытаюсь не отвечать на этот вопрос, ведь ваши соседи снизу утонут в воде, которая сейчас льётся из экрана.
Читать далее - на свой страх и риск. Статья писалась спинным мозгом и глубокой ночью, как следствие я получил натянутую на глобус сущность, которую можно инкапсулировать в технотекст, что бы она вызывала меньше подозрений у случайного читателя. Ну вы поняли уровень, верно?

Сразу хочу оговориться, что в статье не будет ничего принципиально нового - только гайд по паре уже разработанных за нас либ и совсем небольшой mvp код (60+60 строк) как итог.

Кому может быть интересно такое читать? - Тем, кто писал мне в чат хабра и телегу вопросы вроде "А как запустить нейросеть?" или "Как сохранить блокнот колаба?". Мне ни в коем случае не жалко, так что вот маленький туториал на 10 минут вашего экранного времени.

Предлагаю обойтись подобным кратким вступлением и перейти сразу к оглавлению:

• Идея

• Сервер

• Тест сервера (Fast API)

• Клиент

• Тесты

Мы все привыкли, что нейросети интегрируются в коммерческие продукты и работают в облаке. В этой статье же я собираюсь запустить модель локально и использовать не в узкоспециализированном приложении, а совместно с вашим вводом с клавиатуры. Например, сейчас я пишу этот текст и могу нажать сочетание клавиш, которое продолжит его за меня.

Забавно, что вышеупомянутая идея является обращением к тёмному прошлому GPT моделей ввиде подсказок ввода на телефоне, о котором, я полагаю, они хотели бы забыть)

Если компилировать идеи, которые я успел сгенерировать пока думал над вступлением получится примерно это:

1. В стандартных интерфейсах пользователя есть такая сущность, как редактируемые текстовые поля или поля ввода, например: строка ввода пароля от wifi, открытый файл в VScode, адресная строка браузера и т.д.
   И объединяет возможность выделять в них текст, вырезать (ctrl+x), копировать (ctrl+c) и вставлять (ctrl+v). Пока игнорируем специфичные случаи вроде Windows CMD, где эти сочетания передают служебные команды и не выполняют основных функций.

2. ПО может эмулировать ввод пользователя с клавиатуры. Если вы родились с клеймом питониста, то вашими библиотеками до конца жизни будут: PyAutoGui и PyDirectInput.

Размышляя на этим, в памяти всплывают такие бородатые утилиты как PuntoSwitcher и Caramba. Если вы о них никогда не слышали, то вот их принцип работы:

1. Читаем весь поток ввода пользователя с клавиатуры

2. Следим за последним словом, сверяя его со словарями русского и английского

3. Если слово похоже на абракадабру на одном языке, но похоже на другой если его транслитерировать раскладкой (привет=ghbdtn), то пользователь забыл переключить раскладку.

   1. В таком случае переключаем ее сами и стираем мусор, написанный юзером

   2. Пишем, но уже с правильной раскладкой

   3. Делаем предыдущие два пункта очень быстро и незаметно для отвлечённого на клавиатуру пользователя (если он конечно не может в слепую печать)

Повторяем то же упражнение, но уже в контексте работы с нейросетями.

В прошлой статье на тему (см. Реально Бесконечное (лето) RuGPT3.5: Генерация новеллы на ходу нейросетью) могло показаться, что я очень критически отнёсся ко всему семейству лам, но это не так. Просто модель, которая нативно пишет на русском primyerno tak мне совершенно не подходила в конкретном кейсе.

LLama выдающиеся хотя бы потому, что они максимально масштабируемы. Разработчики заранее сделали претрейн всех стандартных размеров, с расчётом локального запуска на GPU любого ценового сегмента.

И раз уж речь пошла про размер модели, то нужно сразу с ним определиться.

Cheat sheet размеров GPT, в зависимости от версии

Если хочется быстро прикинуть, то можно просто умножать кол-во миллиардов нейронов на 2, для получения числа гигов

Выбор сорта кофе лично мне не сильно важен, но раз уж я привык к 50/50 арабике и робусте, а в моей видеокарте менее 15 гигабайт видеопамяти, то остановлюсь я на первой же строке. Вы, в свою очередь, вольны экспериментировать с любыми другими версиями.
В следующей за этой статье из моей серии будет использована уже 13B версия классической ламы, так что не расслабляемся. Там же будут плотные тесты в интересных контекстах.

Запуск модели

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_id = "codellama/CodeLlama-7b-Python-hf"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    use_safetensors=True,
    low_cpu_mem_usage=True,
    device_map="auto",
)

13 гигов VRAM из доступных 15 заполнены - вроде как уложились? Не совсем...
Нам нужна высокая скорость генерации, а она рождается со свободным местом для вычислений. Никто не мешает вам оставить всё как есть, но я добавлю load_in_4bit и загружу модель с уменьшенной битностью - для PoC сойдёт.

from transformers import BitsAndBytesConfig
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_id = "codellama/CodeLlama-7b-Python-hf"
quantization_config = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_compute_dtype=torch.float16
)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=quantization_config,
    use_safetensors=True,
    low_cpu_mem_usage=True,
    device_map="auto",
)

Благо создатели Llama позаботились о нас, и в этот раз разговор с квантованными моделями ограничился парой строк вместо пробирающих до костей мучений.

Для простоты взаимодействия абстрагируемся от токенизатора и сделаем что-то вроде пайплайна в пару строк (да, я казуал)

def txt2txt(text,**kwargs):
  inputs = tokenizer(text, return_tensors="pt").to("cuda")
  output = model.generate(
      inputs["input_ids"],
      **kwargs
  )
  output = output[0].to("cpu")
  return tokenizer.decode(output)

Тут мне резко захотелось разделить функционал сервера (хост с вычислительной мощностью) и клиента(ов) на разные программы. Пока кофеварка пытается пропустить мне пар через кофе вместо воды (ведь я забыл выключить режим капучинатора), мы пишем простой интерфейс взаимодействия с сервером на FastAPI.

from fastapi import FastAP

app = FastAPI()
@app.get("/")
def read_root():
    return {"I am": "alive"}

from pydantic import BaseModel

class Item(BaseModel):
    prompt: str
    add: int
    temperature: float

@app.post("/")
def root(data: Item):
    return {"return": txt2txt(data.prompt,max_new_tokens=data.add,do_sample=True,top_p=0.9,temperature=data.temperature)}

Принимаем POST запросы с промптом, количеством необходимых токенов для генерации и гиперпараметром temperature.

Суть сервера реализована, осталось заварить наше эспрессо. Закидываем вышеизложенный код в файлик server.py и запускаем его таким скриптом:

#with open("server.py","w") as s: s.write(text)
from pycloudflared import try_cloudflare
try_cloudflare(port=8000)
try_cloudflare.terminate(port=8000)
try_cloudflare(port=8000)
!uvicorn server:app --reload

Весь предшествующий код был написан на ванильном python, но запускать я его все равно собирался в google colab. Эта часть нашего нашего зверька (которую я обозвал сервером) может работать как удалённо (на другой машине, у которой лучше GPU), так и локально (на той же, что и клиент). Как вы могли догадаться, Tesla T4 в колабе лучше моей Rx 580, даже есть не брать во внимание кровавые ритуалы жертвоприношения, необходимые для запуска ML проектов на gpu красного вендора.
Как следствие, я решил остановиться на этом.

Устанавливаем requirements пост фактум

!pip install transformers==4.32.1
!pip install fastapi==0.103.0
!pip install bitsandbytes==0.41.1
!pip install accelerate==0.22.0
!pip install "uvicorn[standard]"
!pip install pycloudflared==0.2.0

!pip install transformers==4.32.1
!pip install fastapi==0.103.0
!pip install bitsandbytes==0.41.1
!pip install accelerate==0.22.0
!pip install "uvicorn[standard]"
!pip install pycloudflared==0.2.0

api="""
from transformers import BitsAndBytesConfig
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch

model_id = "codellama/CodeLlama-7b-Python-hf"
quantization_config = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_compute_dtype=torch.float16
)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=quantization_config,
    use_safetensors=True,
    low_cpu_mem_usage=True,
    device_map="auto",
)


def txt2txt(text,**kwargs):
  inputs = tokenizer(text, return_tensors="pt").to("cuda")
  output = model.generate(
      inputs["input_ids"],
      **kwargs
  )
  output = output[0].to("cpu")
  return tokenizer.decode(output)


from fastapi import FastAPI

app = FastAPI()
@app.get("/")
def read_root():
    return {"I am": "alive"}

from pydantic import BaseModel

class Item(BaseModel):
    prompt: str
    add: int
    temperature: float

@app.post("/")
def root(data: Item):
    return {"return": txt2txt(data.prompt,max_new_tokens=data.add,do_sample=True,top_p=0.9,temperature=data.temperature)}
"""
with open("server.py","w") as s: s.write(api)
from pycloudflared import try_cloudflare
try_cloudflare(port=8000)
try_cloudflare.terminate(port=8000)
try_cloudflare(port=8000)
!uvicorn server:app --reload

Тестим

Переходим по последней ссылке и получаем json ответ {"I am":"alive"}.

Добавляем /docs к url и попадаем в автоматически сгенерированный webui

Можем поиграться и протестировать наш бекэнд без клиента.

Пробуем закинуть простой промпт
prompt = 12345

Ответ:
return = 123456789012345

Вроде работает...

Клиент

Ну, тут всё как нельзя проще. Подогреваем молоко паром, превращая его в пенку, параллельно с этим получаем ответы от удалённого сервера и используем их, будто бы GPT прямо у нас под рукой.

import urllib.parse
import requests
import json
import pyautogui as gui
import clipboard
import time

server='https://donations-tunes-institutions-fed.trycloudflare.com/' # Актуальная ссылка

# Или 127.0.0.1 если сервер запущен на том же компьютере что и клиент

def generate(text):
    url = server
    myobj = {'prompt': text,"add":10,"temperature":0.1}
    x = requests.post(url, json = myobj)
    print(x.text)
    return json.loads(x.text)["return"][4:]

Контекст парсер

Используем pyautogui для эмуляции нажатий. Сначала предлагаю взглянуть на кусок кода, а уже после получить оправдания объяснения к нему.

def scan():
    back = clipboard.paste()
    gui.keyDown('shift')
    gui.press('pgup')
    gui.keyUp('shift')
    gui.keyDown('ctrl')
    gui.press('x')
    gui.keyUp('ctrl')
    gui.keyDown('ctrl')
    gui.press('v')
    gui.keyUp('ctrl')
    pasted = clipboard.paste()
    clipboard.copy(back)
    return pasted

Вот такой страшный бойлерплейт, да. В адекватном виде план выглядит так:

1. Shift + PageUP (выделит весь предшествующий контекст)

2. Ctrl+X, Ctrl+V (Вырежет и обратно вставит этот контекст так, что каретка вернётся на место, а выделение спадёт)

3. Берём полученный промпт из буфера и восстанавливаем старый буфер, каким он был до всех наших манипуляций

Конечно, только если вы фанат латте. Можно было бы приготовить раф или капучино, используя в качестве сдвига назад не PageUP, а Home или спам RightArrow. Это добавило бы совместимости, ведь интерфейсов которые не выполняют специфичные команды на правую стрелочку сильно больше чем игнорирующих PageUP.
В сухом остатке: Для изменения способа захвата контекста достаточно изменить сценарий нажатий в функции scan на любой другой.

Теперь выполняем это только тогда, когда нажато некое сочетание клавиш (в моем случае просто F7)

import keyboard
while True:
    time.sleep(0.01)
    try:
        if keyboard.is_pressed("f7"):
            pr=True
        else:
            pr=False
    except:
        pr=False
    if pr:
        old=scan().replace("\r","")
        print([old])
        new=generate(old)
        back = clipboard.paste()
        clipboard.copy(new[len(old):])
        gui.keyDown('ctrl')
        gui.press('v')
        gui.keyUp('ctrl')
        clipboard.copy(back)
        print(new[len(old):])

import urllib.parse
import requests
import json
import pyautogui as gui
import clipboard
import time
server='https://ссылку сюда вставить надо, да/' # Актуальная ссылка

# Или 127.0.0.1 если сервер запущен на том же компьютере что и клиент

def generate(text):
    url = server
    myobj = {'prompt': text,"add":10,"temperature":0.1}
    x = requests.post(url, json = myobj)
    print(x.text)
    return json.loads(x.text)["return"][4:]

def scan():
    back = clipboard.paste()
    gui.keyDown('shift')
    gui.press('pgup')
    gui.keyUp('shift')
    gui.keyDown('ctrl')
    gui.press('x')
    gui.keyUp('ctrl')
    gui.keyDown('ctrl')
    gui.press('v')
    gui.keyUp('ctrl')
    pasted = clipboard.paste()
    clipboard.copy(back)
    return pasted


import keyboard  # using module keyboard
while True:  # making a loop
    time.sleep(0.01)
    try:  # used try so that if user pressed other than the given key error will not be shown
        if keyboard.is_pressed("f7"):
            pr=True
        else:
            pr=False
    except:
        pr=False
    if pr:  # if key 'q' is pressed 
        old=scan().replace("\r","")
        print([old])
        new=generate(old)
        back = clipboard.paste()
        clipboard.copy(new[len(old):])
        gui.keyDown('ctrl')
        gui.press('v')
        gui.keyUp('ctrl')
        clipboard.copy(back)
        print(new[len(old):])

Тесты в полевых условиях

Раз

Два

Три

Conclusion

Отнюдь, я не собираюсь выжимать из себя аналитику и оценку качества работы самой CodeLLama, как это и было изначально постулировано во 2 предложении первого абзаца данного текста.
И этому есть адекватное объяснение. Я не сомневаюсь в качестве самой CodeLLama, но сильно сомневаюсь, что увижу её потенциал, запустив урезанную Ultra lite версию)
Если вам нужен действительно действенный саппорт, можете поставить себе TabNine в VSCode. Даже ссылку оставлять не буду, платные подписки и их реклама - это корпоративное зло.
Вы всегда можете запустить данный небольшой код самостоятельно и проверить его на вашей конкретной задаче. И само собой, мне было бы интересно получить фидбек в комментариях.

Надеюсь, хоть кто-то до этого момента дошёл. Если это был ты, то спасибо тебе, дорогой читатель, осиливший этот небольшой туториал.