Views
No views yet
allenai/olmOCR-2-7B-1025transformers library, please follow the code examples provided on the original model card.1import openai
2import requests
3import base64
4import fitz
5import sys
6from concurrent.futures import ThreadPoolExecutor, as_completed
7
8client = openai.OpenAI(api_key="sk-", base_url="http://ip:port/v1")
9
10model = "winninghealth/olmOCR-2-7B-1025-INT4"
11
12build_no_anchoring_v4_yaml_prompt = "Attached is one page of a document that you must process. Just return the plain text representation of this document as if you were reading it naturally. Convert equations to LateX and tables to HTML.\nIf there are any figures or charts, label them with the following markdown syntax \nReturn your output as markdown."
13
14def render_pdf_to_base64png(doc: fitz.Document, page_num: int, target_longest_image_dim: int = 2048) -> str:
15 page = doc[page_num - 1] # PyMuPDF uses 0-based indexing
16 rect = page.rect
17 width, height = rect.width, rect.height
18 longest_dim = max(width, height)
19
20 # Calculate zoom factor to achieve target dimension
21 zoom = target_longest_image_dim / longest_dim
22
23 # Render page to pixmap
24 mat = fitz.Matrix(zoom, zoom)
25 pix = page.get_pixmap(matrix=mat)
26
27 # Convert pixmap to PNG bytes
28 img_bytes = pix.tobytes("png")
29
30 return base64.b64encode(img_bytes).decode("utf-8")
31
32
33def get_image_base64_from_url(image_url):
34 response = requests.get(image_url)
35 response.raise_for_status()
36 return base64.b64encode(response.content).decode("utf-8")
37
38
39def ocr_page_with_nanonets_s(img_base64):
40 response = client.chat.completions.create(
41 model=model,
42 messages=[
43 {
44 "role": "user",
45 "content": [
46 {
47 "type": "image_url",
48 "image_url": {"url": f"data:image/png;base64,{img_base64}"},
49 },
50 {
51 "type": "text",
52 "text": build_no_anchoring_v4_yaml_prompt,
53 },
54 ],
55 }
56 ],
57 temperature=0.0,
58 max_tokens=15000, # max 16192
59 )
60 return response.choices[0].message.content
61
62
63def process_page(doc, page_num, page_count):
64 img_base64 = render_pdf_to_base64png(doc, page_num, target_longest_image_dim=1288)
65 content = ocr_page_with_nanonets_s(img_base64)
66 return page_num, content
67
68
69# Process all pages concurrently and save to markdown
70if len(sys.argv) < 2:
71 print("Usage: python olmOCR.py <pdf_file_path>")
72 sys.exit(1)
73
74file_path = sys.argv[1]
75output_path = file_path.replace(".pdf", ".md")
76
77# Open PDF once for all operations
78doc = fitz.open(file_path)
79page_count = len(doc)
80
81print(f"Total pages: {page_count}")
82print("Starting OCR processing...\n")
83
84completed_pages = 0
85
86# Open output file for streaming write
87with open(output_path, "w", encoding="utf-8") as f:
88 page_contents = {}
89
90 with ThreadPoolExecutor(max_workers=8) as executor:
91 futures = {
92 executor.submit(process_page, doc, page_num, page_count): page_num for page_num in range(1, page_count + 1)
93 }
94
95 for future in as_completed(futures):
96 page_num, content = future.result()
97 page_contents[page_num] = content
98 completed_pages += 1
99
100 # Display progress
101 progress = (completed_pages / page_count) * 100
102 print(f"Progress: {completed_pages}/{page_count} pages ({progress:.1f}%)")
103
104 # Sort by page number and write to file
105 for i in range(1, page_count + 1):
106 f.write(page_contents[i])
107 # if i < page_count:
108 # f.write("\n\n")
109
110doc.close()
111print(f"\nDone! Output saved to: {output_path}")