Views
No views yet
1from transformers import AutoModelForImageTextToText, AutoProcessor, BitsAndBytesConfig
2from peft import PeftModel
3from PIL import Image
4import torch
5
6# 1. Load Base Model in 4-bit NF4
7model_id = "google/medgemma-1.5-4b-it"
8processor = AutoProcessor.from_pretrained(model_id)
9
10quant_config = BitsAndBytesConfig(
11 load_in_4bit=True,
12 bnb_4bit_quant_type="nf4",
13 bnb_4bit_use_double_quant=True,
14 bnb_4bit_compute_dtype=torch.float16
15)
16
17base_model = AutoModelForImageTextToText.from_pretrained(
18 model_id,
19 device_map="auto",
20 quantization_config=quant_config
21)
22
23# 2. Load this LoRA Adapter
24model = PeftModel.from_pretrained(base_model, "shrish/medgemma-1.5-mm-wsi-module3")
25
26# 3. Format Prompt and Dummy Image
27PROMPT = f"Analyze this 512x512 Bone Marrow Biopsy patch. Does it contain any plasma cells indicative of Multiple Myeloma?"
28messages = [{"role": "user", "content": prompt}]
29formatted_prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
30
31# Note: If feeding actual WSI images in the future, replace the dummy image with the real PIL Image.
32dummy_image = Image.new('RGB', (512, 512), color='black')
33
34inputs = processor(
35 text=formatted_prompt,
36 images=dummy_image,
37 return_tensors="pt",
38 padding=True
39).to(model.device)
40inputs.pop("token_type_ids", None)
41
42# 4. Generate
43with torch.no_grad():
44 outputs = model.generate(**inputs, max_new_tokens=300, do_sample=False)
45
46input_length = inputs["input_ids"].shape[1]
47print(processor.decode(outputs[0, input_length:], skip_special_tokens=True))