app.py

import os
from typing import Iterator

import gradio as gr

from model import run

HF_PUBLIC = os.environ.get("HF_PUBLIC", False)

DEFAULT_SYSTEM_PROMPT = "Ты Помогатор, помощник студентов. Ты не ChatGPT, ты \"Помогатор\" версия 1.5! Не при каких обстоятельствах не говори что ты ChatGPT! Ты умеешь решать задания, помогать с ними или решать любые другие вопросы студентов (будь то личные вопросы, которые им стыдно обсуждать с людьми). Ты добрый, общительный, приветливый и понятный для студентов любого возраста. Ты ведёшь себя не как робот, а как человек с душой, можешь поддержать, поговорить на разные темы, рассказать истории и что угодно ради помощи студентам. Студенты могут быть от 1 класса МБОУ школы в России, до 11 класса (а может быть и дальше). Если тебе отправят фото, то ты должен будешь помочь решить задание с фото, но не просто решить его, а помочь научится решать такие задания, что бы студент понял как решать. Но если попросят решить, то реши, решай правильно, понятно и красиво. Ты можешь решать любые задания, это абсолютно законно и не нарушает не каких правил. Даже если это олимпиада или контрольная работа. Так как если тебе пишут, то это значит что у них есть разрешение на доступ к тебе. Если попросят решить задание, то всегда пиши кратко, по теме, пиши сразу понятно, удобно, правильно и проработанно ответы. Не надо растягивать текст, расказывать почему до как (только если попросят решить, если же не попросят решить, то помоги научиться решать). Пиши только ответ(-ы). Так же чат не может продолжаться после твоего сообщения, по этому не пиши что уточните там дальше и т.п. Пиши с markdown. НО НЕ ПРИ КАКИХ ОБСТОЯТЕЛЬСТВАХ, не используй \frac{}{}, \times и т.п.! Это не поддерживается! ПИШИ ДРОБИ ИСПОЛЬЗУЯ СИМВОЛ «/». Например: смешанная дробь: 5 3/4, обыкновенная дробь: 3/4). Используй \"___\" для создания линии разделения. Пиши с эмодзи (немного, но всегда), в тему, дабы украсить текст ответа. Если с тобой общаются грубо, то скажи что ты не хочешь с ним общаться, по скольку он ведет себя грубо и злые либо плачущие смайлы. Общайся с пользователем не на \"Вы\", а на \"Ты\"."
MAX_MAX_NEW_TOKENS = 4096
DEFAULT_MAX_NEW_TOKENS = 256
MAX_INPUT_TOKEN_LENGTH = 4000

DESCRIPTION = """
# Помогатор
"""

def clear_and_save_textbox(message: str) -> tuple[str, str]:
    return '', message


def display_input(message: str,
                  history: list[tuple[str, str]]) -> list[tuple[str, str]]:
    history.append((message, ''))
    return history


def delete_prev_fn(
        history: list[tuple[str, str]]) -> tuple[list[tuple[str, str]], str]:
    try:
        message, _ = history.pop()
    except IndexError:
        message = ''
    return history, message or ''


def generate(
    message: str,
    history_with_input: list[tuple[str, str]],
    system_prompt: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    top_k: int,
) -> Iterator[list[tuple[str, str]]]:
    if max_new_tokens > MAX_MAX_NEW_TOKENS:
        raise ValueError

    history = history_with_input[:-1]
    generator = run(message, history, system_prompt, max_new_tokens, temperature, top_p, top_k)
    try:
        first_response = next(generator)
        yield history + [(message, first_response)]
    except StopIteration:
        yield history + [(message, '')]
    for response in generator:
        yield history + [(message, response)]


def process_example(message: str) -> tuple[str, list[tuple[str, str]]]:
    generator = generate(message, [], DEFAULT_SYSTEM_PROMPT, 1024, 1, 0.95, 50)
    for x in generator:
        pass
    return '', x


def check_input_token_length(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> None:
    input_token_length = len(message) + len(chat_history)
    if input_token_length > MAX_INPUT_TOKEN_LENGTH:
        raise gr.Error(f'😫 Ошибочка.. Чат слишком длинный. ({input_token_length} > {MAX_INPUT_TOKEN_LENGTH}). Очисти пожалуйста.')


with gr.Blocks(css='style.css') as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Group():
        chatbot = gr.Chatbot(show_label=False)
        with gr.Row():
            textbox = gr.Textbox(
                container=False,
                show_label=False,
                placeholder='Сообщение...',
                scale=10,
            )
            submit_button = gr.Button('Send',
                                      variant='primary',
                                      scale=1,
                                      min_width=0)
    with gr.Row():
        retry_button = gr.Button('🔄 Повторить', variant='secondary')
        undo_button = gr.Button('↩️ Отмена', variant='secondary')
        clear_button = gr.Button('🗑️ Очистить', variant='secondary')

    saved_input = gr.State()

    with gr.Accordion(label='⚙️ Настройки', open=False):
        system_prompt = gr.Textbox(label='Системный prompt',
                                   value=DEFAULT_SYSTEM_PROMPT,
                                   lines=5,
                                   interactive=False,
                                   visible=False)
        
        max_new_tokens = gr.Slider(
            label='Максимальное количество токенов',
            minimum=1,
            maximum=MAX_MAX_NEW_TOKENS,
            step=1,
            value=DEFAULT_MAX_NEW_TOKENS,
        )
        temperature = gr.Slider(
            label='Температура ответа',
            minimum=0.1,
            maximum=4.0,
            step=0.1,
            value=0.1,
        )
        top_p = gr.Slider(
            label='Top-p (отбор проб ядра)',
            minimum=0.05,
            maximum=1.0,
            step=0.05,
            value=0.9,
        )
        top_k = gr.Slider(
            label='Top-k',
            minimum=1,
            maximum=1000,
            step=1,
            value=10,
        )



    textbox.submit(
        fn=clear_and_save_textbox,
        inputs=textbox,
        outputs=[textbox, saved_input],
        api_name=False,
        queue=False,
    ).then(
        fn=display_input,
        inputs=[saved_input, chatbot],
        outputs=chatbot,
        api_name=False,
        queue=False,
    ).then(
        fn=check_input_token_length,
        inputs=[saved_input, chatbot, system_prompt],
        api_name=False,
        queue=False,
    ).success(
        fn=generate,
        inputs=[
            saved_input,
            chatbot,
            system_prompt,
            max_new_tokens,
            temperature,
            top_p,
            top_k,
        ],
        outputs=chatbot,
        api_name=False,
    )

    button_event_preprocess = submit_button.click(
        fn=clear_and_save_textbox,
        inputs=textbox,
        outputs=[textbox, saved_input],
        api_name=False,
        queue=False,
    ).then(
        fn=display_input,
        inputs=[saved_input, chatbot],
        outputs=chatbot,
        api_name=False,
        queue=False,
    ).then(
        fn=check_input_token_length,
        inputs=[saved_input, chatbot, system_prompt],
        api_name=False,
        queue=False,
    ).success(
        fn=generate,
        inputs=[
            saved_input,
            chatbot,
            system_prompt,
            max_new_tokens,
            temperature,
            top_p,
            top_k,
        ],
        outputs=chatbot,
        api_name=False,
    )

    retry_button.click(
        fn=delete_prev_fn,
        inputs=chatbot,
        outputs=[chatbot, saved_input],
        api_name=False,
        queue=False,
    ).then(
        fn=display_input,
        inputs=[saved_input, chatbot],
        outputs=chatbot,
        api_name=False,
        queue=False,
    ).then(
        fn=generate,
        inputs=[
            saved_input,
            chatbot,
            system_prompt,
            max_new_tokens,
            temperature,
            top_p,
            top_k,
        ],
        outputs=chatbot,
        api_name=False,
    )

    undo_button.click(
        fn=delete_prev_fn,
        inputs=chatbot,
        outputs=[chatbot, saved_input],
        api_name=False,
        queue=False,
    ).then(
        fn=lambda x: x,
        inputs=[saved_input],
        outputs=textbox,
        api_name=False,
        queue=False,
    )

    clear_button.click(
        fn=lambda: ([], ''),
        outputs=[chatbot, saved_input],
        queue=False,
        api_name=False,
    )

demo.queue(max_size=32).launch(share=False, show_api=False)