import gradio as gr
from transformers import AutoTokenizer
from optimum.intel import OVModelForCausalLM
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning, message="__array__ implementation doesn't accept a copy keyword")


# 模型與標記器載入（你的原始代碼）
model_id = "hsuwill000/DeepSeek-R1-Distill-Qwen-1.5B-openvino"
print("Loading model...")
model = OVModelForCausalLM.from_pretrained(model_id, device_map="auto")
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)

def respond(prompt, history):
    messages = [
        {"role": "system", "content": "使用中文。"},
        {"role": "user", "content": prompt}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=4096,
        temperature=0.7,
        top_p=0.9,
        do_sample=True
    )
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    response = response.replace("<think>", "**THINK**").replace("</think>", "**THINK**").strip()
    return response

with gr.Blocks() as demo:
    gr.Markdown("# DeepSeek-R1-Distill-Qwen-1.5B-openvino")
    with gr.Tabs():
        with gr.TabItem("聊天"):
            chat_if = gr.Interface(
                fn=respond,
                inputs=gr.Textbox(label="Prompt", placeholder="請輸入訊息..."),
                outputs=gr.Textbox(label="Response", interactive=False),
                api_name="hchat",
                title="DeepSeek-R1-Distill-Qwen-1.5B-openvino",
                description="回傳輸入內容的測試 API",
            )


if __name__ == "__main__":
    print("Launching Gradio app...")
    #demo.queue(api_open=True, max_size=1)
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True, api_open=True, max_size=1)