mwauranjorogekelvin commited on
Commit
f1f87f1
Β·
verified Β·
1 Parent(s): 1010cae

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -75
app.py CHANGED
@@ -5,8 +5,6 @@ import base64
5
  import secrets
6
  import logging
7
  import numpy as np
8
- from hashlib import md5
9
- from collections import OrderedDict
10
  from contextlib import asynccontextmanager
11
 
12
  from fastapi import FastAPI, HTTPException, Security, Depends
@@ -30,48 +28,12 @@ MODEL_PATH = "./model/model.onnx"
30
 
31
  ORT_INTRA_THREADS = int(os.getenv("ORT_INTRA_THREADS", "1"))
32
  ORT_INTER_THREADS = int(os.getenv("ORT_INTER_THREADS", "1"))
33
- CACHE_MAX_SIZE = int(os.getenv("CACHE_MAX_SIZE", "500"))
34
 
35
  # ── Torch global optimizations ────────────────────────────────────────────────
36
  torch.set_grad_enabled(False) # no autograd overhead on tensor ops
37
  torch.set_num_threads(1) # don't compete with ORT threads
38
  torch.set_num_interop_threads(1) # no inter-op parallelism from torch
39
 
40
- # ── LRU Cache ─────────────────────────────────────────────────────────────────
41
- class LRUCache:
42
- def __init__(self, max_size: int):
43
- self.max_size = max_size
44
- self._cache: OrderedDict = OrderedDict()
45
- self.hits = 0
46
- self.misses = 0
47
-
48
- def get(self, key: str):
49
- if key not in self._cache:
50
- self.misses += 1
51
- return None
52
- self._cache.move_to_end(key)
53
- self.hits += 1
54
- return self._cache[key]
55
-
56
- def set(self, key: str, value: str):
57
- if key in self._cache:
58
- self._cache.move_to_end(key)
59
- else:
60
- if len(self._cache) >= self.max_size:
61
- self._cache.popitem(last=False)
62
- self._cache[key] = value
63
-
64
- @property
65
- def size(self):
66
- return len(self._cache)
67
-
68
- @property
69
- def hit_rate(self):
70
- total = self.hits + self.misses
71
- return round(self.hits / total * 100, 1) if total else 0.0
72
-
73
- cache = LRUCache(max_size=CACHE_MAX_SIZE)
74
-
75
  # ── API Key Auth ──────────────────────────────────────────────────────────────
76
  api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
77
 
@@ -90,21 +52,20 @@ opts.inter_op_num_threads = ORT_INTER_THREADS
90
  opts.execution_mode = rt.ExecutionMode.ORT_SEQUENTIAL
91
  opts.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL
92
  opts.optimized_model_filepath = MODEL_PATH + ".opt"
93
- opts.enable_mem_pattern = True # reuse memory allocations across runs
94
- opts.enable_cpu_mem_arena = True # pre-allocate memory pool for ORT
95
 
96
  session = rt.InferenceSession(
97
  MODEL_PATH,
98
  sess_options=opts,
99
  providers=["CPUExecutionProvider"]
100
  )
101
- input_name = session.get_inputs()[0].name # cache once at startup
102
  logger.info("βœ… ONNX model ready")
103
 
104
  # ── Inference ─────────────────────────────────────────────────────────────────
105
  def preprocess(image: Image.Image) -> np.ndarray:
106
- """Normalize and transpose in one contiguous block β€” no extra memory copies."""
107
- image = image.convert("RGB").resize(IMG_SIZE, Image.BILINEAR) # faster than BICUBIC
108
  x = np.ascontiguousarray(image, dtype=np.float32)
109
  x = (x / 255.0 - 0.5) / 0.5
110
  return x.transpose(2, 0, 1)[np.newaxis, :] # [1, 3, H, W]
@@ -113,29 +74,13 @@ def preprocess(image: Image.Image) -> np.ndarray:
113
  def solve_image(image: Image.Image) -> str:
114
  x = preprocess(image)
115
  logits = session.run(None, {input_name: x})[0]
116
- # tokenizer uses torch Tensor ops internally β€” convert only at this boundary
117
  probs = torch.tensor(logits).softmax(-1)
118
  preds, _ = tokenizer.decode(probs)
119
  return preds[0]
120
 
121
 
122
- def solve_with_cache(raw_b64: str, image: Image.Image) -> tuple[str, bool]:
123
- key = md5(raw_b64.encode()).hexdigest()
124
- cached = cache.get(key)
125
- if cached is not None:
126
- logger.info(f"⚑ Cache hit β€” '{cached}'")
127
- return cached, True
128
- text = solve_image(image).strip()[:5]
129
- cache.set(key, text)
130
- return text, False
131
-
132
-
133
  # ── Warmup ────────────────────────────────────────────────────────────────────
134
  def warmup():
135
- """
136
- Run 3 dummy inferences at startup so ORT JIT-compiles the graph
137
- before any real request arrives. Every worker runs this independently.
138
- """
139
  logger.info("Warming up model...")
140
  dummy = Image.new("RGB", IMG_SIZE, color=(128, 128, 128))
141
  for _ in range(3):
@@ -158,7 +103,6 @@ class SolveResponse(BaseModel):
158
  success: bool
159
  text: str = ""
160
  processing_time: float = 0.0
161
- cached: bool = False
162
  error: str = ""
163
 
164
 
@@ -183,13 +127,6 @@ def health():
183
  "quantized": True,
184
  "workers": os.getenv("WEB_CONCURRENCY", "1"),
185
  "intra_threads": ORT_INTRA_THREADS,
186
- "cache": {
187
- "size": cache.size,
188
- "max_size": cache.max_size,
189
- "hits": cache.hits,
190
- "misses": cache.misses,
191
- "hit_rate": f"{cache.hit_rate}%",
192
- },
193
  }
194
 
195
 
@@ -202,17 +139,12 @@ def solve(req: SolveRequest, _: str = Depends(verify_key)):
202
  raw = raw.split(",", 1)[1]
203
 
204
  image = Image.open(io.BytesIO(base64.b64decode(raw)))
205
- text, hit = solve_with_cache(raw, image)
206
 
207
  elapsed = time.time() - start
208
- logger.info(f"βœ… Solved: '{text}' in {elapsed:.3f}s (cached={hit})")
209
 
210
- return SolveResponse(
211
- success=True,
212
- text=text,
213
- processing_time=elapsed,
214
- cached=hit,
215
- )
216
 
217
  except Exception as e:
218
  logger.error(f"Error: {e}")
 
5
  import secrets
6
  import logging
7
  import numpy as np
 
 
8
  from contextlib import asynccontextmanager
9
 
10
  from fastapi import FastAPI, HTTPException, Security, Depends
 
28
 
29
  ORT_INTRA_THREADS = int(os.getenv("ORT_INTRA_THREADS", "1"))
30
  ORT_INTER_THREADS = int(os.getenv("ORT_INTER_THREADS", "1"))
 
31
 
32
  # ── Torch global optimizations ────────────────────────────────────────────────
33
  torch.set_grad_enabled(False) # no autograd overhead on tensor ops
34
  torch.set_num_threads(1) # don't compete with ORT threads
35
  torch.set_num_interop_threads(1) # no inter-op parallelism from torch
36
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
37
  # ── API Key Auth ──────────────────────────────────────────────────────────────
38
  api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
39
 
 
52
  opts.execution_mode = rt.ExecutionMode.ORT_SEQUENTIAL
53
  opts.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL
54
  opts.optimized_model_filepath = MODEL_PATH + ".opt"
55
+ opts.enable_mem_pattern = True
56
+ opts.enable_cpu_mem_arena = True
57
 
58
  session = rt.InferenceSession(
59
  MODEL_PATH,
60
  sess_options=opts,
61
  providers=["CPUExecutionProvider"]
62
  )
63
+ input_name = session.get_inputs()[0].name
64
  logger.info("βœ… ONNX model ready")
65
 
66
  # ── Inference ─────────────────────────────────────────────────────────────────
67
  def preprocess(image: Image.Image) -> np.ndarray:
68
+ image = image.convert("RGB").resize(IMG_SIZE, Image.BILINEAR)
 
69
  x = np.ascontiguousarray(image, dtype=np.float32)
70
  x = (x / 255.0 - 0.5) / 0.5
71
  return x.transpose(2, 0, 1)[np.newaxis, :] # [1, 3, H, W]
 
74
  def solve_image(image: Image.Image) -> str:
75
  x = preprocess(image)
76
  logits = session.run(None, {input_name: x})[0]
 
77
  probs = torch.tensor(logits).softmax(-1)
78
  preds, _ = tokenizer.decode(probs)
79
  return preds[0]
80
 
81
 
 
 
 
 
 
 
 
 
 
 
 
82
  # ── Warmup ────────────────────────────────────────────────────────────────────
83
  def warmup():
 
 
 
 
84
  logger.info("Warming up model...")
85
  dummy = Image.new("RGB", IMG_SIZE, color=(128, 128, 128))
86
  for _ in range(3):
 
103
  success: bool
104
  text: str = ""
105
  processing_time: float = 0.0
 
106
  error: str = ""
107
 
108
 
 
127
  "quantized": True,
128
  "workers": os.getenv("WEB_CONCURRENCY", "1"),
129
  "intra_threads": ORT_INTRA_THREADS,
 
 
 
 
 
 
 
130
  }
131
 
132
 
 
139
  raw = raw.split(",", 1)[1]
140
 
141
  image = Image.open(io.BytesIO(base64.b64decode(raw)))
142
+ text = solve_image(image).strip()[:5]
143
 
144
  elapsed = time.time() - start
145
+ logger.info(f"βœ… Solved: '{text}' in {elapsed:.3f}s")
146
 
147
+ return SolveResponse(success=True, text=text, processing_time=elapsed)
 
 
 
 
 
148
 
149
  except Exception as e:
150
  logger.error(f"Error: {e}")