import huggingface_hub as hf_hub
import time
import openvino_genai as ov_genai
import numpy as np
import gradio as gr
import re

# 下載模型
model_ids = [
    "OpenVINO/Qwen3-0.6B-int4-ov",
    "OpenVINO/Qwen3-0.6B-int8-ov",
    "OpenVINO/Qwen3-1.7B-int4-ov",
    "OpenVINO/Qwen3-1.7B-int8-ov"
]

model_name_to_id = {
    "Qwen3-0.6B-int4-ov": "OpenVINO/Qwen3-0.6B-int4-ov",
    "Qwen3-0.6B-int8-ov": "OpenVINO/Qwen3-0.6B-int8-ov",
    "Qwen3-1.7B-int4-ov": "OpenVINO/Qwen3-1.7B-int4-ov",
    "Qwen3-1.7B-int8-ov": "OpenVINO/Qwen3-1.7B-int8-ov"
}

for model_id in model_ids:
    model_path = model_id.split("/")[-1]  # Extract model name
    try:
      hf_hub.snapshot_download(model_id, local_dir=model_path, local_dir_use_symlinks=False)
      print(f"Successfully downloaded {model_id} to {model_path}") # Optional: Print confirmation
    except Exception as e:
      print(f"Error downloading {model_id}: {e}") # Handle download errors gracefully

# 建立推理管線 (Initialize with a default model first)
device = "CPU"
default_model_name = "Qwen3-0.6B-int4-ov"  # Choose a default model
model_path = model_name_to_id[default_model_name].split("/")[-1]

def generate_response(prompt, model_name):
    global pipe, tokenizer  # Access the global variables

    # Check if the model needs to be changed
    model_id = model_name_to_id[model_name]
    new_model_path = model_id.split("/")[-1]

    
    print(f"Switching to model: {model_name}")
    pipe = ov_genai.LLMPipeline(new_model_path, device)
    tokenizer = pipe.get_tokenizer()
    tokenizer.set_chat_template(tokenizer.chat_template)


    try:
        generated = pipe.generate([prompt], max_length=1024)
        tokenpersec=f'{generated.perf_metrics.get_throughput().mean:.2f}'

        return tokenpersec, generated
    except Exception as e:
        return "發生錯誤", "發生錯誤", f"生成回應時發生錯誤：{e}"
    

# 建立 Gradio 介面
model_choices = list(model_name_to_id.keys())

demo = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(lines=5, label="輸入提示 (Prompt)"),
        gr.Dropdown(choices=model_choices, value=default_model_name, label="選擇模型") # Added dropdown
    ],
    outputs=[
        gr.Textbox(label="tokens/sec"),
        gr.Textbox(label="回應")
    ],
    title="Qwen3 Model Inference",
    description="基於 Qwen3 推理應用，支援思考過程分離與 GUI。"
)

if __name__ == "__main__":
    demo.launch()