import torch
import gradio as gr
import spaces
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
from threading import Thread

model_id = "llm-jp/llm-jp-3-8x1.8b-instruct3"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    device_map="auto", 
    torch_dtype=torch.bfloat16
)

@spaces.GPU
def generate_text(system_prompt, user_input, max_length=512, temperature=0.7, top_p=0.95):
    chat = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_input},
    ]
    
    input_ids = tokenizer.apply_chat_template(
        chat, 
        add_generation_prompt=True, 
        return_tensors="pt"
    ).to(model.device)
    
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    
    generate_kwargs = {
        "input_ids": input_ids,
        "streamer": streamer,
        "max_new_tokens": max_length,
        "do_sample": True,
        "temperature": temperature,
        "top_p": top_p,
        "repetition_penalty": 1.05
    }
    
    thread = Thread(target=model.generate, kwargs=generate_kwargs)
    thread.start()
    
    response = ""
    for text in streamer:
        response += text
    
    return response

with gr.Blocks() as demo:
    gr.Markdown("# LLM-JP-3-8x1.8b-instruct3 非公式デモ")
    gr.Markdown("国立情報学研究所大規模言語モデル研究開発センターの開発した日本語大規模言語モデル LLM-JP-3-8x1.8b-instruct3 の非公式デモ。詳細は[こちらの記事](https://llm-jp.nii.ac.jp/blog/2025/03/27/moe3.html)をご覧ください。推論時に Hugging Face の ZeroGPU(A100) を使用しています。")
    
    with gr.Row():
        with gr.Column():
            system_prompt = gr.Textbox(
                label="システムプロンプト", 
                value="以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。",
                lines=2
            )
            user_input = gr.Textbox(label="プロンプト", lines=5, placeholder="プロンプトを入力してください")
            
            with gr.Row():
                max_length = gr.Slider(label="最大トークン数", minimum=10, maximum=1024, value=512, step=1)
                temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, value=0.7, step=0.1)
                top_p = gr.Slider(label="Top-p", minimum=0.1, maximum=1.0, value=0.95, step=0.05)
            
            submit_btn = gr.Button("生成")
        
        with gr.Column():
            output = gr.Textbox(label="生成結果", lines=20)
    
    submit_btn.click(
        fn=generate_text,
        inputs=[system_prompt, user_input, max_length, temperature, top_p],
        outputs=output
    )
    
    gr.Examples(
        examples=[
            ["以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。", "自然言語処理とは何か"],
            ["以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。", "日本の四季について教えてください"],
            ["あなたは優秀な物語作家です。", "猫と犬が友達になる短い物語を書いてください。"]
        ],
        inputs=[system_prompt, user_input]
    )

demo.launch()