zixianma02 commited on
Commit
733d337
·
verified ·
1 Parent(s): 018ebe4

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +4 -2
README.md CHANGED
@@ -59,7 +59,8 @@ checkpoint_dir = "allenai/MolmoWeb-4B"
59
  model = AutoModelForImageTextToText.from_pretrained(
60
  checkpoint_dir,
61
  trust_remote_code=True,
62
- dtype="auto",
 
63
  device_map="auto",
64
  )
65
 
@@ -122,7 +123,8 @@ inputs = processor.apply_chat_template(
122
  padding=True,
123
  )
124
 
125
- inputs = {k: v.to("cuda") for k, v in inputs.items()}
 
126
 
127
  with torch.inference_mode():
128
  output = model.generate(**inputs, max_new_tokens=200)
 
59
  model = AutoModelForImageTextToText.from_pretrained(
60
  checkpoint_dir,
61
  trust_remote_code=True,
62
+ torch_dtype=torch.float32, # we recommend using the default float32 precision
63
+ attn_implementation="sdpa",
64
  device_map="auto",
65
  )
66
 
 
123
  padding=True,
124
  )
125
 
126
+ # Remove token_type_ids: HF uses it to enable bidirectional attention for image tokens; molmoweb is trained with causal attention only
127
+ inputs = {k: v.to("cuda") for k, v in inputs.items() if k != "token_type_ids"}
128
 
129
  with torch.inference_mode():
130
  output = model.generate(**inputs, max_new_tokens=200)