import os
import gradio as gr
from transformers import pipeline

# 環境変数からAPIトークンを取得
hf_api_token = os.environ.get("HF_API_TOKEN")

# トークンが設定されていない場合のエラー処理
if not hf_api_token:
    raise ValueError("Hugging FaceのAPIトークンが設定されていません。環境変数 'HF_API_TOKEN' を設定してください。")

# Hugging Faceのパイプラインを初期化
chatbot_pipeline = pipeline(
    "text-generation",
    model="Sakalti/Tara-3.8B-v1.1",
    use_auth_token=hf_api_token  # 環境変数から取得したトークンを使用
)

# 応答を生成する関数
def respond(prompt, system_message, max_new_tokens, temperature, top_p, top_k):
    """
    チャットボットの応答を生成する。
    """
    # プロンプトをシステムメッセージと結合
    full_prompt = f"{system_message}\nユーザー: {prompt}\nAI:"
    
    # パイプラインで応答を生成
    response = chatbot_pipeline(
        full_prompt,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        num_return_sequences=1
    )
    # 応答テキストを返す
    return response[0]["generated_text"]

# Gradioでチャットボットのインターフェースを構築
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="あなたはフレンドリーなチャットボットです。", label="システムメッセージ"),
        gr.Slider(minimum=1, maximum=2048, value=768, step=1, label="新規トークン最大"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="温度"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (核 sampling)"),
        gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-k"),
    ],
    concurrency_limit=30  # 同時リクエスト数の上限
)

# アプリケーションを起動
demo.launch()