Views
No views yet

1from transformers import AutoProcessor, AutoModelForImageTextToText
2from PIL import Image
3import requests
4import torch
5from jinja2 import Template
6
7checkpoint_dir = "allenai/MolmoWeb-8B"
8
9model = AutoModelForImageTextToText.from_pretrained(
10 checkpoint_dir,
11 trust_remote_code=True,
12 torch_dtype=torch.float32, # we recommend using the default float32 precision
13 attn_implementation="sdpa",
14 device_map="auto",
15)
16
17processor = AutoProcessor.from_pretrained(
18 checkpoint_dir,
19 trust_remote_code=True,
20 padding_side="left",
21)
22
23
24MOLMOWEB_THINK_TEMPLATE = Template(
25"""
26# GOAL
27{{ task_description }}
28
29# PREVIOUS STEPS
30{% for action in past_actions: -%}
31## Step {{ action['index'] }}
32THOUGHT: {{ action['thought'] }}
33ACTION: {{ action['action'] }}
34{% endfor %}
35# CURRENTLY ACTIVE PAGE
36Page {{ page_index }}: {{ page_title }} | {{ page_url }}
37
38# NEXT STEP
39
40"""
41)
42
43task_description = "Tell me about the Ai2 PIROR team's recent projects"
44past_actions = []
45user_message = MOLMOWEB_THINK_TEMPLATE.render(
46 page_title=None,
47 page_url="about:blank",
48 page_index=0,
49 task_description=task_description,
50 past_actions=[]
51)
52system_message = "molmo_web_think"
53prompt = f"{system_message}: {user_message}"
54
55blank_image = Image.new("RGB", (1280, 720), color="white")
56
57image_messages = [
58 {
59 "role": "user",
60 "content": [
61 {"type": "text", "text": prompt},
62 {"type": "image", "image": blank_image},
63 ]
64 }
65]
66
67inputs = processor.apply_chat_template(
68 image_messages,
69 tokenize=True,
70 add_generation_prompt=True,
71 return_tensors="pt",
72 return_dict=True,
73 padding=True,
74)
75
76# Remove token_type_ids: HF uses it to enable bidirectional attention for image tokens; molmoweb is trained with causal attention only
77inputs = {k: v.to("cuda") for k, v in inputs.items() if k != "token_type_ids"}
78
79with torch.inference_mode():
80 output = model.generate(**inputs, max_new_tokens=200)
81
82generated_tokens = output[0, inputs["input_ids"].size(1):]
83print(processor.decode(generated_tokens, skip_special_tokens=True))