# 安装依赖
# !pip install transformers accelerate bitsandbytes huggingface_hub

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSeq2SeqLM, AutoProcessor, AutoModelForVision2Seq

import torch

# 加载模型（使用 4-bit 量化直接加载，避免 OOM）
model_name = "ByteDance-Seed/UI-TARS-1.5-7B"  # 或 UI-TARS-1.5-7B（如果你有权访问）

model = AutoModelForVision2Seq.from_pretrained(
    model_name,                           
    device_map="auto",         # ⬅️ 4-bit 量化
    torch_dtype=torch.float16,
    quantization_config={
        "load_in_4bit": True,
        "bnb_4bit_quant_type": "nf4",        # ✅ 必须是 nf4（CPU 只支持这个）
        "bnb_4bit_compute_dtype": torch.float16,
        "bnb_4bit_use_double_quant": True,   # 可选：减少 0.4% 体积
    },
    low_cpu_mem_usage=True
)

tokenizer = AutoTokenizer.from_pretrained(model_name)

# # 保存量化后的模型到本地（或 Hugging Face）
# model.save_pretrained("./ui-tars-8b-4bit")
# tokenizer.save_pretrained("./ui-tars-8b-4bit")

def greet(input):

    # prepare the model input
    prompt = "Give me a short introduction to large language model."
    prompt = input
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=True,  # Switches between thinking and non-thinking modes. Default is True.
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

    # conduct text completion
    generated_ids = model.generate(**model_inputs, max_new_tokens=32768)
    output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist()

    # parsing thinking content
    try:
        # rindex finding 151668 (</think>)
        index = len(output_ids) - output_ids[::-1].index(151668)
    except ValueError:
        index = 0

    thinking_content = tokenizer.decode(
        output_ids[:index], skip_special_tokens=True
    ).strip("\n")
    content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")

    # print("thinking content:", thinking_content)
    # print("content:", content)

    return "thinking content:" + thinking_content + "\n" + "content:" + content


demo = gr.Interface(fn=greet, inputs="text", outputs="text")
demo.launch()