Views
No views yet


1from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
2from qwen_vl_utils import process_vision_info
3
4# Load the Bpe-vocab-n-OCR model with optimized parameters
5model = Qwen2VLForConditionalGeneration.from_pretrained(
6 "prithivMLmods/Tokenized-OCR", torch_dtype="auto", device_map="auto"
7)
8
9# Recommended acceleration for performance optimization:
10# model = Qwen2VLForConditionalGeneration.from_pretrained(
11# "prithivMLmods/Tokenized-OCR",
12# torch_dtype=torch.bfloat16,
13# attn_implementation="flash_attention_2",
14# device_map="auto",
15# )
16
17# Load the default processor for Bpe-vocab-n-OCR
18processor = AutoProcessor.from_pretrained("prithivMLmods/Tokenized-OCR")
19
20# Define the input messages with both an image and a text prompt
21messages = [
22 {
23 "role": "user",
24 "content": [
25 {
26 "type": "image",
27 "image": "https://flux-generated.com/sample_image.jpeg",
28 },
29 {"type": "text", "text": "Extract and return the tokenized OCR text from the image, ensuring each word is accurately recognized and separated by commas."},
30 ],
31 }
32]
33
34# Prepare the input for inference
35text = processor.apply_chat_template(
36 messages, tokenize=False, add_generation_prompt=True
37)
38image_inputs, video_inputs = process_vision_info(messages)
39inputs = processor(
40 text=[text],
41 images=image_inputs,
42 videos=video_inputs,
43 padding=True,
44 return_tensors="pt",
45)
46inputs = inputs.to("cuda")
47
48# Generate the output
49generated_ids = model.generate(**inputs, max_new_tokens=256)
50generated_ids_trimmed = [
51 out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
52]
53output_text = processor.batch_decode(
54 generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
55)
56print(output_text)