from vllm import LLM, SamplingParams
import gradio as gr

llm = LLM(model="EQUES/JPharmatron-7B")
sampling_params = SamplingParams(temperature=0.0, max_tokens=512)

# 応答関数
def chat_bot(user_input, history):
    # システムプロンプトを付与して文脈を明確にする
    prompt = "以下は親切で何でも答えてくれるAIアシスタントとの会話です。\n"
    for user_msg, bot_msg in history:
        prompt += f"ユーザー: {user_msg}\nアシスタント: {bot_msg}\n"
    prompt += f"ユーザー: {user_input}\nアシスタント:"

    outputs = llm.generate(prompt, sampling_params)
    reply = outputs[0].outputs[0].text.strip()
    history.append((user_input, reply))
    return reply, history

# Gradio UIの構築
with gr.Blocks() as demo:
    gr.Markdown("# 💊 製薬に関する質問をしてみてください。")
    gr.Markdown("※ あくまでデモとしての利用のみに止めてください。")
    
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="あなたのメッセージを入力してください。")
    clear = gr.Button("チャット履歴をクリア")

    state = gr.State([])

    def respond(message, history):
        response, history = chat_bot(message, history)
        return history, history

    msg.submit(respond, [msg, state], [chatbot, state])
    clear.click(lambda: ([], []), None, [chatbot, state])

demo.launch()