This repository contains the ONNX version of the CodeFormula model optimized with JPQD (Joint Pruning, Quantization, and Distillation) quantization for efficient inference.
1import onnxruntime as ort
2import numpy as np
3from PIL import Image
4import cv2
5
6# Load the CodeFormula ONNX model
7model_path = "CodeFormula.onnx"
8session = ort.InferenceSession(model_path)
9
10def preprocess_image(image_path):
11 """Preprocess image for CodeFormula model"""
12 # Load image at 120 DPI as specified in model documentation
13 image = Image.open(image_path).convert('RGB')
14
15 # Resize to appropriate dimensions (adjust based on model requirements)
16 # CodeFormula expects 120 DPI images
17 image = image.resize((800, 600)) # Example dimensions
18
19 # Convert to numpy array
20 image_array = np.array(image)
21
22 # For this example, we'll create a dummy token sequence
23 # In practice, you'd use the actual preprocessing pipeline
24 dummy_input = np.random.randint(0, 50827, (1, 10)).astype(np.int64)
25
26 return dummy_input
27
28def recognize_code_or_formula(image_path):
29 """Recognize code or formula from image"""
30
31 # Preprocess image
32 input_tokens = preprocess_image(image_path)
33
34 # Run inference
35 outputs = session.run(None, {"input": input_tokens})
36 logits = outputs[0] # Shape: [1, 10, 50827]
37
38 # Get predicted tokens (simplified decoding)
39 predicted_tokens = np.argmax(logits[0], axis=-1)
40
41 return predicted_tokens
42
43# Example usage
44image_path = "code_snippet.jpg"
45tokens = recognize_code_or_formula(image_path)
46print(f"Predicted tokens: {tokens}")
1import onnxruntime as ort
2import numpy as np
3from typing import List, Union
4import cv2
5from PIL import Image
6
7class CodeFormulaONNX:
8 """ONNX wrapper for CodeFormula model"""
9
10 def __init__(self, model_path: str = "CodeFormula.onnx"):
11 """Initialize CodeFormula ONNX model"""
12 print(f"Loading CodeFormula model: {model_path}")
13 self.session = ort.InferenceSession(model_path)
14
15 # Get model info
16 self.input_name = self.session.get_inputs()[0].name
17 self.input_shape = self.session.get_inputs()[0].shape
18 self.output_names = [output.name for output in self.session.get_outputs()]
19
20 # Model vocabulary size
21 self.vocab_size = 50827
22
23 print(f"✓ Model loaded successfully")
24 print(f" Input: {self.input_name} {self.input_shape}")
25 print(f" Vocabulary size: {self.vocab_size}")
26
27 def preprocess_image(self, image: Union[str, np.ndarray]) -> np.ndarray:
28 """
29 Preprocess image for CodeFormula inference
30
31 Args:
32 image: Image path or numpy array
33
34 Returns:
35 Input tensor for the model
36 """
37
38 if isinstance(image, str):
39 # Load image from path
40 pil_image = Image.open(image).convert('RGB')
41 image_array = np.array(pil_image)
42 else:
43 image_array = image
44
45 # CodeFormula expects 120 DPI images
46 # Adjust size based on DPI requirements
47 height, width = image_array.shape[:2]
48
49 # Resize to maintain 120 DPI (adjust as needed)
50 target_height, target_width = 600, 800 # Example dimensions
51 if height != target_height or width != target_width:
52 image_array = cv2.resize(image_array, (target_width, target_height))
53
54 # Convert to grayscale for better OCR (optional)
55 if len(image_array.shape) == 3:
56 gray = cv2.cvtColor(image_array, cv2.COLOR_RGB2GRAY)
57 else:
58 gray = image_array
59
60 # Apply image preprocessing for better recognition
61 # Enhance contrast
62 clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
63 enhanced = clahe.apply(gray)
64
65 # For this demonstration, create dummy token input
66 # In practice, you would tokenize the image using the actual preprocessing pipeline
67 dummy_tokens = np.random.randint(0, self.vocab_size, self.input_shape).astype(np.int64)
68
69 return dummy_tokens
70
71 def predict(self, input_tokens: np.ndarray) -> np.ndarray:
72 """Run model prediction"""
73
74 # Validate input shape
75 if input_tokens.shape != tuple(self.input_shape):
76 print(f"Warning: Input shape {input_tokens.shape} != expected {self.input_shape}")
77
78 # Run inference
79 outputs = self.session.run(None, {self.input_name: input_tokens})
80
81 return outputs[0] # Return logits
82
83 def decode_output(self, logits: np.ndarray) -> List[int]:
84 """Decode model output logits to tokens"""
85
86 # Get most likely tokens
87 predicted_tokens = np.argmax(logits[0], axis=-1)
88
89 return predicted_tokens.tolist()
90
91 def recognize(self, image: Union[str, np.ndarray]) -> dict:
92 """
93 Recognize code or formula from image
94
95 Args:
96 image: Image path or numpy array
97
98 Returns:
99 Dictionary with recognition results
100 """
101
102 # Preprocess image
103 input_tokens = self.preprocess_image(image)
104
105 # Run inference
106 logits = self.predict(input_tokens)
107
108 # Decode output
109 predicted_tokens = self.decode_output(logits)
110
111 # Analyze output pattern (simplified)
112 result = {
113 "predicted_tokens": predicted_tokens,
114 "sequence_length": len(predicted_tokens),
115 "max_logit": float(np.max(logits)),
116 "mean_confidence": float(np.mean(np.max(logits[0], axis=-1))),
117 "type": self._classify_output_type(predicted_tokens)
118 }
119
120 return result
121
122 def _classify_output_type(self, tokens: List[int]) -> str:
123 """Classify if output is likely code or formula (simplified heuristic)"""
124
125 # This is a simplified classification
126 # In practice, you'd use the actual tokenizer to decode and analyze
127
128 # Placeholder classification based on token patterns
129 if len(tokens) > 5:
130 return "code"
131 else:
132 return "formula"
133
134 def benchmark(self, num_iterations: int = 100) -> dict:
135 """Benchmark model performance"""
136
137 print(f"Running benchmark with {num_iterations} iterations...")
138
139 # Create dummy input
140 dummy_input = np.random.randint(0, self.vocab_size, self.input_shape).astype(np.int64)
141
142 # Warmup
143 for _ in range(5):
144 _ = self.predict(dummy_input)
145
146 # Benchmark
147 import time
148 times = []
149
150 for i in range(num_iterations):
151 start_time = time.time()
152 _ = self.predict(dummy_input)
153 end_time = time.time()
154 times.append(end_time - start_time)
155
156 if (i + 1) % 10 == 0:
157 print(f" Progress: {i + 1}/{num_iterations}")
158
159 # Calculate statistics
160 times = np.array(times)
161 stats = {
162 "mean_time_ms": float(np.mean(times) * 1000),
163 "std_time_ms": float(np.std(times) * 1000),
164 "min_time_ms": float(np.min(times) * 1000),
165 "max_time_ms": float(np.max(times) * 1000),
166 "median_time_ms": float(np.median(times) * 1000),
167 "throughput_fps": float(1.0 / np.mean(times))
168 }
169
170 return stats
171
172# Example usage
173def main():
174 # Initialize model
175 codeformula = CodeFormulaONNX("CodeFormula.onnx")
176
177 # Example 1: Recognize from image file
178 image_path = "code_example.jpg"
179 try:
180 result = codeformula.recognize(image_path)
181 print(f"Recognition result: {result}")
182 except FileNotFoundError:
183 print("Example image not found, using dummy data...")
184
185 # Example 2: Recognize from numpy array
186 dummy_image = np.random.randint(0, 255, (600, 800, 3), dtype=np.uint8)
187 result = codeformula.recognize(dummy_image)
188 print(f"Dummy recognition result: {result}")
189
190 # Example 3: Performance benchmark
191 print("\nRunning performance benchmark...")
192 stats = codeformula.benchmark(50)
193 print(f"Benchmark results:")
194 print(f" Mean inference time: {stats['mean_time_ms']:.2f} ms")
195 print(f" Throughput: {stats['throughput_fps']:.1f} FPS")
196
197if __name__ == "__main__":
198 main()
1# Note: This is a conceptual example
2# The actual integration would depend on tokenizer availability
3
4from transformers import AutoTokenizer
5import onnxruntime as ort
6
7# If tokenizer is available
8try:
9 tokenizer = AutoTokenizer.from_pretrained("ds4sd/CodeFormula")
10
11 def decode_tokens(token_ids):
12 return tokenizer.decode(token_ids, skip_special_tokens=True)
13
14except:
15 print("Tokenizer not available, using dummy decoding")
16
17 def decode_tokens(token_ids):
18 return f"<decoded_sequence_length_{len(token_ids)}>"
1def process_code_images_batch(image_paths, batch_size=4):
2 """Process multiple code images in batches"""
3
4 codeformula = CodeFormulaONNX("CodeFormula.onnx")
5 results = []
6
7 for i in range(0, len(image_paths), batch_size):
8 batch = image_paths[i:i+batch_size]
9
10 batch_results = []
11 for image_path in batch:
12 result = codeformula.recognize(image_path)
13 batch_results.append({
14 "image_path": image_path,
15 "recognition": result
16 })
17
18 results.extend(batch_results)
19 print(f"Processed batch {i//batch_size + 1}/{(len(image_paths)-1)//batch_size + 1}")
20
21 return results
22
23# Usage
24image_list = ["code1.jpg", "code2.jpg", "formula1.jpg"]
25batch_results = process_code_images_batch(image_list)
1@techreport{Docling,
2 author = {Deep Search Team},
3 month = {8},
4 title = {{Docling Technical Report}},
5 url={https://arxiv.org/abs/2408.09869},
6 eprint={2408.09869},
7 doi = "10.48550/arXiv.2408.09869",
8 version = {1.0.0},
9 year = {2024}
10}
11
12@misc{zhang2022opt,
13 title={OPT: Open Pre-trained Transformer Language Models},
14 author={Susan Zhang and Stephen Roller and Naman Goyal and Mikel Artetxe and Moya Chen and Shuohui Chen and Christopher Dewan and Mona Diab and Xian Li and Xi Victoria Lin and Todor Mihaylov and Myle Ott and Sam Shleifer and Kurt Shuster and Daniel Simig and Punit Singh Koura and Anjali Sridhar and Tianlu Wang and Luke Zettlemoyer},
15 year={2022},
16 eprint={2205.01068},
17 archivePrefix={arXiv},
18 primaryClass={cs.CL}
19}
