File size: 2,401 Bytes
2b0192c fcb5bdc 2b0192c 98dfb19 60603bf 98dfb19 2b0192c 59b0e70 2b0192c 59b0e70 2b0192c 60603bf 2b0192c 59b0e70 2b0192c 59b0e70 620ce13 59b0e70 98dfb19 620ce13 98dfb19 59b0e70 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
# 安装依赖
# !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()
|