A BERT-based sequence classification model that routes computer vision questions to appropriate specialized modules. Classifies questions into 4 task categories: VQA, Captioning, Grounding, and Geometry.
1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2import torch
3
4model_name = "beingamanforever/ICM"
5tokenizer = AutoTokenizer.from_pretrained(model_name)
6model = AutoModelForSequenceClassification.from_pretrained(model_name)
7
8questions = [
9 "What is the distance between the two trees?",
10 "Describe what the child is wearing.",
11 "Is the traffic light green?",
12 "Box the location of the blue umbrella."
13]
14
15inputs = tokenizer(questions, return_tensors="pt", padding=True, truncation=True)
16with torch.no_grad():
17 logits = model(**inputs).logits
18 predictions = torch.argmax(logits, dim=-1)
19
20for q, pred in zip(questions, predictions):
21 print(f"{q} → {model.config.id2label[pred.item()]}")