CodeLama в вашей клавиатуре | Локальный Copilot для любого поля ввода

Стою я значит утром (около 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, в зависимости от версии

Model	Original size	Quantized size
7B	13 GB	3.9 GB
13B	24 GB	7.8 GB
30B	60 GB	19.5 GB
65B	120 GB	38.5 GB

Если хочется быстро прикинуть, то можно просто умножать кол-во миллиардов нейронов на 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 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)}

Принимаем 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

Try cloud...что?

PyCloudFlared - python обертка над утилитой от cloudflare, которая позволяет создать временный HTTP туннель, как NGROK, только без регистрации и api.
Каждое поднятие тоннеля генерирует новый URL вида:
https://X-X-X-X.trycloudflare.com/
Однако, если использовать порт, на котором уже поднимался тоннель то обертка выдаёт уже существующий URL, что является фактической недоработкой.
Ведь тоннель может к тому моменту уже лечь, однако URL будет возвращаться старый. Для этого я, на всякий случай, сразу делаю перезапуск тоннеля и получаю новый, но точно актуальный URL

Весь предшествующий код был написан на ванильном 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

Весь код сервера на ipython

!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):])

Весь код клиента на python

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):])

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

Раз

Два

Три

я дописал in a cycle, скриншот немного outdated

---

Conclusion

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

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