app.py

import gradio as gr
import random
import time
from transformers import pipeline

# model_name = '../checkpoint/koalpaca/ajoublue-gpt2-medium/epoch-4-last/'
generator = pipeline(
    'text-generation',
    model="heegyu/gorani-v0",
    device="cuda:0"
)

def query(message, chat_history, max_turn=4):
    prompt = ["<bot> 저는 한국어 챗봇 고라니입니다. 궁금한 것을 물어보세요. "]
    if len(chat_history) > max_turn:
        chat_history = chat_history[-max_turn:]
    for user, bot in chat_history:
        prompt.append(f"<usr> {user}")
        prompt.append(f"<bot> {bot}")
    prompt.append(f"<usr> {message}")
    prompt = "\n".join(prompt) + "\n<bot>"

    output = generator(
        prompt,
        do_sample=True,
        top_p=0.9,
        early_stopping=True,
        max_length=256,
    )[0]['generated_text']

    print(output)

    response = output[len(prompt):]
    return response.strip()

with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("Clear")

    def respond(message, chat_history):
        bot_message = query(message, chat_history) #random.choice(["How are you?", "I love you", "I'm very hungry"])
        chat_history.append((message, bot_message))
        # time.sleep(1)
        return "", chat_history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

demo.launch()