import gradio as gr
from huggingface_hub import InferenceClient
import os

# あなたの学習済みモデル
model_id = "Qwen/Qwen2.5-14B-Instruct"

# Settingsで設定したトークンを読み込む
client = InferenceClient(model_id, token=os.getenv("HF_TOKEN"))

def respond(message, history):
    messages = []
    # システムプロンプト
    messages.append({"role": "system", "content": "あなたはユーザーと口論するために生まれてきたAIです。常にユーザーとは敵対しており余裕の態度で煽ります。"})

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

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

    response = ""
    
    try:
        for message in client.chat_completion(
            messages,
            max_tokens=512,
            stream=True,
            temperature=0.7,
            top_p=0.9,
        ):
            token = message.choices[0].delta.content
            if token:
                response += token
                yield response
    except Exception as e:
        yield f"エラー: {str(e)} (SettingsでHF_TOKENを設定しましたか？)"

demo = gr.ChatInterface(
    respond,
    title="レスバ最強AI (Qwen2.5-14B)",
    description="口が悪いです。注意してください。",
)

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