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

# 下載模型
model_ids = [
    "OpenVINO/Qwen3-0.6B-int4-ov",
    "OpenVINO/Qwen3-1.7B-int4-ov",
    #"OpenVINO/Qwen3-4B-int4-ov",#不可用
    "OpenVINO/Qwen3-8B-int4-ov",
    "OpenVINO/Qwen3-14B-int4-ov",
    
]

model_name_to_full_id = {model_id.split("/")[-1]: model_id for model_id in model_ids}  #Create Dictionary

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"
InUsed_model_name = "Qwen3-0.6B-int4-ov"  # Choose a default model
pipe = ov_genai.LLMPipeline(InUsed_model_name, device)
tokenizer = pipe.get_tokenizer()
tokenizer.set_chat_template(tokenizer.chat_template)
config = ov_genai.GenerationConfig()
config.max_new_tokens = 1024

def streamer(subword):
    print(subword, end='', flush=True)
    return False
    
def generate_response(prompt, model_name):
    global pipe, tokenizer, InUsed_model_name
    if(InUsed_model_name!=model_name):
        model_path = model_name
        del pipe
        gc.collect()
        print(f"Switching to model: {model_name}")
        pipe = ov_genai.LLMPipeline(model_path, device)
        tokenizer = pipe.get_tokenizer()
        tokenizer.set_chat_template(tokenizer.chat_template)
        InUsed_model_name=model_name

    try:
        generated = pipe.generate([prompt], config, streamer)
        tokenpersec=f'{generated.perf_metrics.get_throughput().mean:.2f}'
        print(f"\nModel:{model_name} TPS:{tokenpersec}\n")
        return tokenpersec, generated
    except Exception as e:
        return "發生錯誤", "發生錯誤", f"生成回應時發生錯誤：{e}"
    

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

demo = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(lines=5, label="輸入提示 (Prompt)"),
        gr.Dropdown(choices=model_choices, value=InUsed_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()