app.py

import gradio as gr
import requests
import json
import threading
import os
import datetime
from requests.exceptions import RequestException

stop_generation = threading.Event()
API_URL = os.environ.get('API_URL')
API_KEY = os.environ.get('API_KEY')

headers = {
    "Authorization": f"Bearer {API_KEY}",
    "Content-Type": "application/json"
}

session = requests.Session()

DEFAULT_PARAMS = {
    "temperature": 0.8,
    "top_p": 0.95,
    "top_k": 40,
    "frequency_penalty": 0,
    "presence_penalty": 0,
    "repetition_penalty": 1.1,
    "max_tokens": 512
}

def get_timestamp():
    return datetime.datetime.now().strftime("%H:%M:%S")

def predict(message, history, system_prompt, temperature, top_p, top_k, frequency_penalty, presence_penalty, repetition_penalty, max_tokens):
    global stop_generation, session
    stop_generation.clear()  

    history_format = [{"role": "system", "content": system_prompt}]
    for human, assistant in history:
        history_format.append({"role": "user", "content": human})
        if assistant:
            history_format.append({"role": "assistant", "content": assistant})
    history_format.append({"role": "user", "content": message})

    if not message.startswith(('*', '"')):
        print(f"<|system|> {system_prompt}")
        print(f"{get_timestamp()} <|user|> {message}")

    current_params = {
        "temperature": temperature,
        "top_p": top_p,
        "top_k": top_k,
        "frequency_penalty": frequency_penalty,
        "presence_penalty": presence_penalty,
        "repetition_penalty": repetition_penalty,
        "max_tokens": max_tokens
    }

    non_default_params = {k: v for k, v in current_params.items() if v != DEFAULT_PARAMS[k]}
    
    if non_default_params and not message.startswith(('*', '"')):
        for param, value in non_default_params.items():
            print(f"{param}={value}")

    data = {
        "model": "meta-llama/Meta-Llama-3.1-405B-Instruct",
        "messages": history_format,
        "stream": True,
        "temperature": temperature,
        "top_p": top_p,
        "top_k": top_k,
        "frequency_penalty": frequency_penalty,
        "presence_penalty": presence_penalty,
        "repetition_penalty": repetition_penalty,
        "max_tokens": max_tokens
    }

    try:
        with session.post(API_URL, headers=headers, data=json.dumps(data), stream=True) as response:
            partial_message = ""
            for line in response.iter_lines():
                if stop_generation.is_set():
                    response.close()
                    break
                if line:
                    line = line.decode('utf-8')
                    if line.startswith("data: "):
                        if line.strip() == "data: [DONE]":
                            break
                        try:
                            json_data = json.loads(line[6:])
                            if 'choices' in json_data and json_data['choices']:
                                content = json_data['choices'][0]['delta'].get('content', '')
                                if content:
                                    partial_message += content
                                    yield partial_message
                        except json.JSONDecodeError:
                            continue  

        if partial_message:
            yield partial_message

    except RequestException as e:
        print(f"Request error: {e}")
        yield f"An error occurred: {str(e)}"

def import_chat(custom_format_string):
    try:
        sections = custom_format_string.split('<|')

        imported_history = []
        system_prompt = ""

        for section in sections:
            if section.startswith('system|>'):
                system_prompt = section.replace('system|>', '').strip()
            elif section.startswith('user|>'):
                user_message = section.replace('user|>', '').strip()
                imported_history.append([user_message, None])
            elif section.startswith('assistant|>'):
                assistant_message = section.replace('assistant|>', '').strip()
                if imported_history:
                    imported_history[-1][1] = assistant_message
                else:
                    imported_history.append(["", assistant_message])

        return imported_history, system_prompt
    except Exception as e:
        print(f"Error importing chat: {e}")
        return None, None

def export_chat(history, system_prompt):
    export_data = f"<|system|> {system_prompt}\n\n"
    if history is not None:
        for user_msg, assistant_msg in history:
            export_data += f"<|user|> {user_msg}\n\n"
            if assistant_msg:
                export_data += f"<|assistant|> {assistant_msg}\n\n"
    return export_data

def stop_generation_func():
    global stop_generation, session
    stop_generation.set()
    session.close()  
    session = requests.Session()  

with gr.Blocks(theme='gradio/monochrome') as demo:
    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(value=[])
            msg = gr.Textbox(label="Message")
            with gr.Row():
                clear = gr.Button("Clear")
                regenerate = gr.Button("Regenerate")
                stop_btn = gr.Button("Stop")
            with gr.Row():
                with gr.Column(scale=4):
                    import_textbox = gr.Textbox(label="Import textbox", lines=5)
                with gr.Column(scale=1):
                    export_button = gr.Button("Export Chat")
                    import_button = gr.Button("Import Chat")

        with gr.Column(scale=1):
            system_prompt = gr.Textbox("", label="System Prompt", lines=5)
            temperature = gr.Slider(0, 2, value=0.8, step=0.01, label="Temperature")
            top_p = gr.Slider(0, 1, value=0.95, step=0.01, label="Top P")
            top_k = gr.Slider(1, 500, value=40, step=1, label="Top K")
            frequency_penalty = gr.Slider(-2, 2, value=0, step=0.1, label="Frequency Penalty")
            presence_penalty = gr.Slider(-2, 2, value=0, step=0.1, label="Presence Penalty")
            repetition_penalty = gr.Slider(0.01, 5, value=1.1, step=0.01, label="Repetition Penalty")
            max_tokens = gr.Slider(1, 4096, value=512, step=1, label="Max Output (max_tokens)")

    def user(user_message, history):
        history = history or []
        return "", history + [[user_message, None]]

    def bot(history, system_prompt, temperature, top_p, top_k, frequency_penalty, presence_penalty, repetition_penalty, max_tokens):
        global stop_generation
        history = history or []
        if not history:
            return history
        user_message = history[-1][0]
        bot_message = predict(user_message, history[:-1], system_prompt, temperature, top_p, top_k, frequency_penalty, presence_penalty, repetition_penalty, max_tokens)
        history[-1][1] = ""
        for chunk in bot_message:
            if stop_generation.is_set():
                history[-1][1] += " [Generation stopped]"
                break
            history[-1][1] = chunk
            yield history
        stop_generation.clear()

    def regenerate_response(history, system_prompt, temperature, top_p, top_k, frequency_penalty, presence_penalty, repetition_penalty, max_tokens):
        if history and len(history) > 0:
            last_user_message = history[-1][0]
            history[-1][1] = None  
            for new_history in bot(history, system_prompt, temperature, top_p, top_k, frequency_penalty, presence_penalty, repetition_penalty, max_tokens):
                yield new_history
        else:
            yield []

    def import_chat_wrapper(custom_format_string):
        imported_history, imported_system_prompt = import_chat(custom_format_string)
        return imported_history, imported_system_prompt

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, [chatbot, system_prompt, temperature, top_p, top_k, frequency_penalty, presence_penalty, repetition_penalty, max_tokens], chatbot
    )

    clear.click(lambda: None, None, chatbot, queue=False)

    regenerate.click(
        regenerate_response,
        [chatbot, system_prompt, temperature, top_p, top_k, frequency_penalty, presence_penalty, repetition_penalty, max_tokens],
        chatbot
    )

    import_button.click(import_chat_wrapper, inputs=[import_textbox], outputs=[chatbot, system_prompt])

    export_button.click(
        export_chat,
        inputs=[chatbot, system_prompt],
        outputs=[import_textbox]
    )

    stop_btn.click(stop_generation_func, inputs=[], outputs=[])

if __name__ == "__main__":
    demo.launch(debug=True)