from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig

MODEL_NAME = "IlyaGusev/saiga_13b_lora"
DEFAULT_MESSAGE_TEMPLATE = "<s>{role}\n{content}</s>\n"
DEFAULT_SYSTEM_PROMPT = "Ты — Сайга, русскоязычный автоматический ассистент. Ты разговариваешь с людьми и помогаешь им."

class Conversation:
    def __init__(
        self,
        message_template=DEFAULT_MESSAGE_TEMPLATE,
        system_prompt=DEFAULT_SYSTEM_PROMPT,
        start_token_id=1,
        bot_token_id=9225
    ):
        self.message_template = message_template
        self.start_token_id = start_token_id
        self.bot_token_id = bot_token_id
        self.messages = [{
            "role": "system",
            "content": system_prompt
        }]

    def get_start_token_id(self):
        return self.start_token_id

    def get_bot_token_id(self):
        return self.bot_token_id

    def add_user_message(self, message):
        self.messages.append({
            "role": "user",
            "content": message
        })

    def add_bot_message(self, message):
        self.messages.append({
            "role": "bot",
            "content": message
        })

    def get_prompt(self, tokenizer):
        final_text = ""
        for message in self.messages:
            message_text = self.message_template.format(**message)
            final_text += message_text
        final_text += tokenizer.decode([self.start_token_id, self.bot_token_id])
        return final_text.strip()

def generate(model, tokenizer, prompt, generation_config):
    data = tokenizer(prompt, return_tensors="pt")
    data = {k: v.to(model.device) for k, v in data.items()}
    output_ids = model.generate(
        **data,
        generation_config=generation_config
    )[0]
    output_ids = output_ids[len(data["input_ids"][0]):]
    output = tokenizer.decode(output_ids, skip_special_tokens=True)
    return output.strip()

config = PeftConfig.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    config.base_model_name_or_path,
    load_in_8bit=True,
    torch_dtype=torch.float16,
    device_map="auto"
)
model = PeftModel.from_pretrained(
    model,
    MODEL_NAME,
    torch_dtype=torch.float16
)
model.eval()

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
generation_config = GenerationConfig.from_pretrained(MODEL_NAME)
print(generation_config)

inputs = ["Почему трава зеленая?", "Сочини длинный рассказ, обязательно упоминая следующие объекты. Дано: Таня, мяч"]
for inp in inputs:
    conversation = Conversation()
    conversation.add_user_message(inp)
    prompt = conversation.get_prompt(tokenizer)

    output = generate(model, tokenizer, prompt, generation_config)
    print(inp)
    print(output)
    print()
    print("==============================")
    print()

Examples:

User: Почему трава зеленая? 
Saiga: Зеленый цвет у растений обусловлен наличием хлорофилла в их листьях. Хлорофилл - это пигмент, который отвечает за фотосинтез, то есть процесс преобразования солнечной энергии в химическую энергию. Он находится в специальных органелах, называемых хлоропластами, которые имеются только у растительных клеток. Когда свет попадает на листья, хлорофилл переходит из состояния "растворенного" в "состояние связанное", что позволяет ему использовать энергию света для синтеза сахаров. Этот процесс происходит в хлоропластах и приводит к образованию кислорода и углекислого газа. В результате этого процесса листья становятся зелеными.

User: Сочини длинный рассказ, обязательно упоминая следующие объекты. Дано: Таня, мяч
Saiga: Татьяна была очень увлеченной девочкой, которая всегда любила играть в футбол. Она была настоящей звездой на поле, и ее команда всегда побеждала благодаря ее таланту и энергии.
Однажды, когда Татьяна играла в футбол со своими друзьями, она случайно потеряла свой мяч. Она была очень расстроена, так как это был ее любимый мяч, который она использовала уже несколько лет.
Но несмотря на то, что она была расстроена, Татьяна решила продолжить игру. Она начала искать свой мяч, но ничего не нашла. В конце концов, она решила вернуться домой и попросить помощи у своих родителей.
Когда Татьяна пришла домой, она рассказала своим родителям о том, что произошло. Они были очень сочувствующими и предложили ей найти новый мяч.
Татьяна была очень благодарна своим родителям за их помощь, и они отправились вместе в магазин спортивного инвентаря. Там они выбрали для нее новый мяч, который был почти такой же, как старый.
После того, как Татьяна получила новый мяч, она вернулась на поле и продолжила играть. Но теперь она знала, что ее любовь к игре не зависит от одного мяча - она могла играть и без него.
Также, Татьяна поняла, что ее родители всегда будут рядом ей, чтобы помочь ей в трудные моменты. И она была очень благодарна им за это.
Таким образом, Татьяна научилась жить без своего старого мяча, и она стала более сильной и уверенной в себе после этого случая. Ее любовь к игре только укрепилась, и она продолжала играть в футбол каждую свободную минуту.

User: Реши уравнение 4x + 5 = 21
Saiga: Для решения данного уравнения необходимо выполнить следующие шаги:
1. Перенесите все члены в левую часть уравнения на другую сторону равенства: 
4x + 5 = 21
2. Отмените оба члена, переведя их в правую часть уравнения:
4x = 21 - 5
3. Разделите оба члена на общее деление:
4x / 4 = (21 - 5) / 4
4x / 4 = 5 - 5/4
4x / 4 = 0
4x = 0
4x = 0 * 4
4x = 0

Ответ: x = 0.

User: Заполни пропуски в предложении. Дано: Я пытался ____ от маньяка, но он меня настиг
Saiga: Я пытался убежать от маньяка, но он меня настиг.

v2:

revision fe7034727ce126498559df277b08da76001edc16
wandb link
7 datasets: ru_turbo_alpaca, ru_turbo_saiga, ru_sharegpt_cleaned, oasst1_ru_main_branch, ru_turbo_alpaca_evol_instruct (1/2 iteration), ru_instruct_gpt4
Datasets merging script: create_chat_set.py
saiga13b_v2 vs saiga13b_v1: 105-10-61

v1:

revision bd4823198d443ab0a18be3768c9985de6e249d43
wandb link
4 datasets: ru_turbo_alpaca, ru_turbo_saiga, ru_sharegpt_cleaned, oasst1_ru_main_branch
Datasets merging script: create_chat_set.py
Loss: 0.870
Context length: 2000
Conversational template: "<s>{role}\n{content}</s>"
Possible roles: ["system", "user", "bot"]
System prompt: "Ты — Сайга, русскоязычный автоматический ассистент. Ты разговариваешь с людьми и помогаешь им."