Views
No views yet
1import torch
2import torch.nn as nn
3from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
4from safetensors.torch import load_file
5from PIL import Image
6
7# Define classifier
8class VLMClassifier(nn.Module):
9 def __init__(self, base_model, num_labels=2):
10 super().__init__()
11 self.base_model = base_model
12 hidden_size = base_model.config.hidden_size
13 self.classifier = nn.Sequential(
14 nn.Linear(hidden_size, 512),
15 nn.ReLU(),
16 nn.Dropout(0.1),
17 nn.Linear(512, num_labels)
18 )
19
20 def forward(self, input_ids, attention_mask, pixel_values, image_grid_thw):
21 outputs = self.base_model(
22 input_ids=input_ids,
23 attention_mask=attention_mask,
24 pixel_values=pixel_values,
25 image_grid_thw=image_grid_thw,
26 output_hidden_states=True,
27 return_dict=True
28 )
29 hidden_states = outputs.hidden_states[-1]
30 pooled = hidden_states.mean(dim=1)
31 return self.classifier(pooled.float())
32
33# Load model
34processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct")
35base_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
36 "Qwen/Qwen2.5-VL-3B-Instruct",
37 torch_dtype=torch.bfloat16,
38 device_map="cuda:0"
39)
40model = VLMClassifier(base_model, num_labels=2).to("cuda")
41model.classifier.load_state_dict(load_file("classifier_head.safetensors"))
42model.eval()
43
44# Inference
45img = Image.open("table.png").convert("RGB")
46messages = [[{
47 "role": "user",
48 "content": [
49 {"type": "image", "image": img},
50 {"type": "text", "text": "Classify this table."}
51 ]
52}]]
53texts = [processor.apply_chat_template(m, tokenize=False, add_generation_prompt=True) for m in messages]
54inputs = processor(text=texts, images=[img], padding=True, return_tensors="pt")
55
56with torch.no_grad():
57 logits = model(
58 inputs["input_ids"].to("cuda"),
59 inputs["attention_mask"].to("cuda"),
60 inputs["pixel_values"].to("cuda", dtype=torch.bfloat16),
61 inputs["image_grid_thw"].to("cuda")
62 )
63 prob_sct = torch.softmax(logits, dim=-1)[0, 1].item()
64
65print(f"P(SCT) = {prob_sct:.3f}")
66# Use threshold 0.3 for fewer false negatives
67is_sct = prob_sct >= 0.3classifier_head.safetensors - Classifier head weightsclassifier_config.json - Model configurationconfig.json - Base model confignotebooks/ - Training and testing notebooks