Update README.md

README.md

@@ -10,15 +10,19 @@ pip install git+https://github.com/tic-top/transformers.git
 from transformers import AutoModelForVision2Seq, AutoProcessor
 from PIL import Image
 import torch
-device = "cuda:0"
-repo = "kirp/kosmos2_5"
-dtype = torch.bfloat16
-model = AutoModelForVision2Seq.from_pretrained(repo, device_map = device).to(dtype)
+device = "cuda:2"
+repo = "kosmos2_5"
+dtype = torch.float16
+# dtype = torch.bfloat16
+model = AutoModelForVision2Seq.from_pretrained(repo, device_map = device, torch_dtype=dtype)
+# print(model)
+# exit(0)
 processor = AutoProcessor.from_pretrained(repo)
 
 path = "receipt_00008.png"
 image = Image.open(path)
-prompt = "<ocr>" # "<md>"
+prompt = "<ocr>" 
+# prompt = "<md>"
 inputs = processor(text=prompt, images=image, return_tensors="pt", max_patches=4096)
 
 raw_width, raw_height = image.size
@@ -29,8 +33,8 @@ scale_width = raw_width / width
 inputs = {k: v.to(device) if v is not None else None for k, v in inputs.items()}
 inputs["flattened_patches"] = inputs["flattened_patches"].to(dtype)
 with torch.no_grad():
-    generated_text = model.generate(**inputs, max_new_tokens=256)
-    
+    generated_text = model.generate(**inputs, max_length=4096)
+generated_text = processor.batch_decode(generated_text)
 import re, os
 def postprocess(y, scale_height, scale_width, result_path=None):
     y = (
@@ -40,6 +44,7 @@ def postprocess(y, scale_height, scale_width, result_path=None):
         .replace("</image>", "")
         .replace(prompt, "")
     )
+    print(y)
     pattern = r"<bbox><x_\d+><y_\d+><x_\d+><y_\d+></bbox>"
     bboxs_raw = re.findall(pattern, y)
     lines = re.split(pattern, y)[1:]
@@ -67,5 +72,5 @@ def postprocess(y, scale_height, scale_width, result_path=None):
     else:
         print(info)
 
-postprocess(processor.batch_decode(generated_text)[0],scale_height, scale_width)
+postprocess(generated_text[0], scale_height, scale_width)
 ```