import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import warnings
warnings.filterwarnings("ignore")

"""
Sarashinaモデルを使用したGradioチャットボット
Hugging Face Transformersライブラリを使用してローカルでモデルを実行
"""

# モデルとトークナイザーの初期化
MODEL_NAME = "sbintuitions/sarashina2.2-3b-instruct-v0.1"

print("モデルを読み込み中...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto" if torch.cuda.is_available() else None,
    trust_remote_code=True
)
print("モデルの読み込みが完了しました。")

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    """
    チャットボットの応答を生成する関数
    """
    try:
        # システムメッセージと会話履歴を含むプロンプトを構築
        conversation = ""
        if system_message.strip():
            conversation += f"システム: {system_message}\n"
        
        # 会話履歴を追加
        for user_msg, bot_msg in history:
            if user_msg:
                conversation += f"ユーザー: {user_msg}\n"
            if bot_msg:
                conversation += f"アシスタント: {bot_msg}\n"
        
        # 現在のメッセージを追加
        conversation += f"ユーザー: {message}\nアシスタント: "
        
        # トークン化
        inputs = tokenizer.encode(conversation, return_tensors="pt")
        
        # GPU使用時はCUDAに移動
        if torch.cuda.is_available():
            inputs = inputs.cuda()
        
        # 応答生成
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_new_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
                repetition_penalty=1.1
            )
        
        # 生成されたテキストをデコード
        generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # 応答部分のみを抽出
        full_response = generated[len(conversation):].strip()
        
        # 不要な部分を除去
        if "ユーザー:" in full_response:
            full_response = full_response.split("ユーザー:")[0].strip()
        
        return full_response
            
    except Exception as e:
        return f"エラーが発生しました: {str(e)}"

# Gradio Blocksを使用したカスタムチャットインターフェース
with gr.Blocks(
    title="🤖 Sarashina Chatbot",
    theme=gr.themes.Soft()
) as demo:
    
    gr.Markdown("# 🤖 Sarashina Chatbot")
    gr.Markdown("Sarashina2.2-3b-instruct モデルを使用した日本語チャットボットです。")
    
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(height=500)
            msg = gr.Textbox(
                label="メッセージを入力してください",
                placeholder="こんにちは！何かお手伝いできることはありますか？",
                lines=2
            )
            clear = gr.Button("会話をクリア")
        
        with gr.Column(scale=1):
            gr.Markdown("### 設定")
            system_message = gr.Textbox(
                value="あなたは親切で知識豊富な日本語アシスタントです。ユーザーの質問に丁寧に答えてください。",
                label="システムメッセージ",
                lines=3
            )
            max_tokens = gr.Slider(
                minimum=1,
                maximum=1024,
                value=512,
                step=1,
                label="最大新規トークン数"
            )
            temperature = gr.Slider(
                minimum=0.1,
                maximum=2.0,
                value=0.7,
                step=0.1,
                label="Temperature (創造性)"
            )
            top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.05,
                label="Top-p (多様性制御)"
            )
    
    # 例文
    gr.Examples(
        examples=[
            ["こんにちは！今日はどんなことを話しましょうか？"],
            ["日本の文化について教えてください。"],
            ["簡単なレシピを教えてもらえますか？"],
            ["プログラミングについて質問があります。"],
        ],
        inputs=msg,
        label="例文"
    )
    
    def user(message, history):
        return "", history + [[message, None]]
    
    def bot(history, system_message, max_tokens, temperature, top_p):
        history[-1][1] = respond(
            history[-1][0],
            history[:-1],
            system_message,
            max_tokens,
            temperature,
            top_p
        )
        return history
    
    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, [chatbot, system_message, max_tokens, temperature, top_p], chatbot
    )
    
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,  # パブリックリンクを作成
        show_api=True,
        debug=True
    )