Update README.md

README.md

@@ -39,36 +39,22 @@ from uform.gen_model import VLMForCausalLM, VLMProcessor
 model = VLMForCausalLM.from_pretrained("unum-cloud/uform-gen-chat")
 processor = VLMProcessor.from_pretrained("unum-cloud/uform-gen-chat")
 
-messages = [
-    {"role": "system", "content": "You are a helpful assistant."},
-    {"role": "user", "content": "<image> {Your message}"}
-]
-
-image = processor.image_processor(Image.open("zebra.jpg")).unsqueeze(0)
-
-input_ids = processor.tokenizer.apply_chat_template(
-    messages, return_tensors="pt", add_generation_prompt=True
-)
-
-attention_mask = torch.ones(1, input_ids.shape[1] + processor.num_image_latents - 1)
-
-inputs = {
-    "input_ids": input_ids,
-    "attention_mask": attention_mask,
-    "images": image,
-}
-
-outputs = model.generate(
-    **inputs,
-    do_sample=False,
-    use_cache=True,
-    max_new_tokens=1024,
-    eos_token_id=32001,
-    pad_token_id=processor.tokenizer.pad_token_id,
-)
-
-message = processor.batch_decode(outputs[:, inputs["input_ids"].shape[1]:-1])
-
+prompt = "What do you see?"
+image = Image.open("zebra.jpg")
+
+inputs = processor(texts=[prompt], images=[image], return_tensors="pt")
+with torch.inference_mode():
+     output = model.generate(
+        **inputs,
+        do_sample=False,
+        use_cache=True,
+        max_new_tokens=128,
+        eos_token_id=32001,
+        pad_token_id=processor.tokenizer.pad_token_id
+    )
+
+prompt_len = inputs["input_ids"].shape[1]
+decoded_text = processor.batch_decode(output[:, prompt_len:])[0]
 ```