Ubuntu中部署谷歌的Medgemma X光影像识别模型
1、创建虚拟环境:
conda create -n medgemma python=3.10 -y && conda activate medgemma
2、克隆项目
cd medgemma
pip install -U huggingface_hub -i https://pypi.tuna.tsinghua.edu.cn/simple
(模型下载前需要先去huggingface上找到谷歌的这个项目获取授权)
https://huggingface.co/google/medgemma-4b-it
5、编写测试代码app.py:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
from transformers import AutoProcessor, AutoModelForImageTextToText
from PIL import Image
import torch
# ———— (可选)显式开启 TorchDynamo ————
# torch._dynamo.enable()
# 模型标识
model_id = "google/medgemma-4b-it"
# 加载模型(不禁用 Dynamo)
model = AutoModelForImageTextToText.from_pretrained(
model_id,
torch_dtype=torch.bfloat16,
device_map="auto",
)
# 加载预处理器
processor = AutoProcessor.from_pretrained(model_id)
# 读取本地图片
image_path = "chest.jpg"
image = Image.open(image_path).convert("RGB")
# 构造对话消息
messages = [
{
"role": "system",
"content": [{"type": "text", "text": "You are an expert radiologist."}]
},
{
"role": "user",
"content": [
{"type": "text", "text": "Describe this X-ray"},
{"type": "image", "image": image}
]
}
]
# 处理为模型输入
inputs = processor.apply_chat_template(
messages,
add_generation_prompt=True,
tokenize=True,
return_dict=True,
return_tensors="pt"
).to(model.device, dtype=torch.bfloat16)
input_len = inputs["input_ids"].shape[-1]
# 生成回答
with torch.inference_mode():
outputs = model.generate(
**inputs,
max_new_tokens=500,
do_sample=False,
use_cache=True,
)
# 切分 & 解码
gen_tokens = outputs[0][input_len:]
decoded = processor.decode(gen_tokens, skip_special_tokens=True)
# 打印
print(decoded)
6、运行程序(将需要检查的x光图片放在项目文件下,命名为chest.jpg)
python app.py
评论
发表评论