Views
No views yet

1from datasets import load_dataset
2from transformers import AutoModelForCausalLM, AutoProcessor
3import torch
4
5device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
6model = AutoModelForCausalLM.from_pretrained("mynkchaudhry/Florence-2-FT-DocVQA").to(device)
7processor = AutoProcessor.from_pretrained("mynkchaudhry/Florence-2-FT-DocVQA")
8
9data = load_dataset("HuggingFaceM4/DocumentVQA")
10
11def run_example(task_prompt, text_input, image):
12 prompt = task_prompt + text_input
13 if image.mode != "RGB":
14 image = image.convert("RGB")
15 inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
16 generated_ids = model.generate(
17 input_ids=inputs["input_ids"],
18 pixel_values=inputs["pixel_values"],
19 max_new_tokens=1024,
20 num_beams=3
21 )
22 generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
23 parsed_answer = processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
24 return parsed_answer
25
26for idx in range(3):
27 print(run_example("DocVQA", 'What do you see in this image?', data['train'][idx]['image']))