import gradio as gr
import spaces

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

MODEL_NAME = "cognitivecomputations/dolphin-2.9.3-mistral-nemo-12b"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
    low_cpu_mem_usage=True
)
print("Model loaded successfully!")

@spaces.GPU
def chat(message, history):
    # historyがtupleリストかdictリストか両対応にする
    messages = []
    if history:
        for item in history:
            if isinstance(item, dict):
                # dict形式（role/content）の場合
                messages.append(item)
            elif isinstance(item, (list, tuple)) and len(item) == 2:
                # tuple形式（user, assistant）の場合
                user, assistant = item
                if user:
                    messages.append({"role": "user", "content": user})
                if assistant:
                    messages.append({"role": "assistant", "content": assistant})
            else:
                # 予期せぬ形式ならスキップ or エラー処理
                continue

    # 今のメッセージを追加
    messages.append({"role": "user", "content": message})

    # プロンプト作成（dolphinモデルなのでシンプルに）
    full_prompt = ""
    for msg in messages:
        if msg["role"] == "user":
            full_prompt += f"ユーザー: {msg['content']}\n"
        elif msg["role"] == "assistant":
            full_prompt += f"アシスタント: {msg['content']}\n\n"
    full_prompt += "アシスタント:"

    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        try:
            outputs = model.generate(
                **inputs,
                max_new_tokens=128,
                do_sample=True,
                temperature=0.7,
                top_p=0.9,
                repetition_penalty=1.1,
                pad_token_id=tokenizer.eos_token_id
            )
            response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
        except Exception as e:
            response = f"ごめんね、エラー出ちゃった… ({str(e)}) もう一回試してみて～！"

    # 返却はGradioが期待するdict形式で追加
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": response})

    return "", history

with gr.Blocks(title="Uncensored 12B Chat") as demo:
    gr.Markdown("# Dolphin 12B Uncensored Chat")
    gr.Markdown("軽くて速いよ～！なんでも話してね……検閲なしだぜ！！！")

    chatbot = gr.Chatbot(height=600)
    msg = gr.Textbox(placeholder="メッセージ入力してEnter...")
    clear = gr.Button("クリア")

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

demo.launch()