Update README.md

README.md

@@ -45,9 +45,10 @@ pip install perceptron
 ### Usage
 
 ```python
+import torch
 from transformers import AutoModelForCausalLM, AutoProcessor
-from transformers.utils.import_utils import is_torch_cuda_available
 from transformers.image_utils import load_image
+from transformers.utils.import_utils import is_torch_cuda_available
 
 def document_to_messages(document: list[dict]):
     messages, images = [], []
@@ -62,14 +63,18 @@ def document_to_messages(document: list[dict]):
             messages.append({"role": role, "content": content})
     return messages, images
 
+hf_path = "PerceptronAI/Isaac-0.2-2B-Preview"
+device, dtype = ("cuda",torch.bfloat16) if is_torch_cuda_available() else ("cpu",torch.float32)
+
 # Load model/processor from the checkpoint
-checkpoint_path = "PerceptronAI/Isaac-0.2-2B-Preview"
-processor = AutoProcessor.from_pretrained(checkpoint_path, trust_remote_code=True)
-device, dtype = ("cuda","bfloat16") if is_torch_cuda_available() else ("cpu","float32")
+processor = AutoProcessor.from_pretrained(hf_path, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
-    checkpoint_path, trust_remote_code=True, vision_attn_implementation="flash_attention_2", dtype = dtype
-).to(device=device)
+    hf_path, trust_remote_code=True, vision_attn_implementation="flash_attention_2"
+)
+model = model.to(device=device, dtype=dtype)
+model.eval()
 
+# Prepare input for generation
 document = [
     {
         "type": "text",
@@ -87,20 +92,21 @@ document = [
         "role": "user",
     },
 ]
-
-# Prepare inputs for generation
 messages, images = document_to_messages(document)
-text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+text = processor.apply_chat_template(
+    messages, tokenize=False, add_generation_prompt=True
+)
 inputs = processor(text=text, images=images, return_tensors="pt")
 
-# Generation
+# Generate text using the model
 generated_ids = model.generate(
-    tensor_stream=inputs["tensor_stream"].to(device),
+    tensor_stream=inputs["tensor_stream"].to(next(model.parameters()).device),
     max_new_tokens=256,
     do_sample=False,
 )
-generated_text = processor.tokenizer.decode(generated_ids[0], skip_special_tokens=False)
-print(f"\n Output: {generated_text}")
-
+generated_text = processor.tokenizer.decode(
+    generated_ids[0], skip_special_tokens=False
+)
+print(f"\nFull generated output:\n{generated_text}")
 ```