Spaces:
Paused
Paused
Suvadeep Das
commited on
Update app.py
Browse files
app.py
CHANGED
@@ -21,9 +21,8 @@ if HF_TOKEN:
|
|
21 |
_model = None
|
22 |
_tokenizer = None
|
23 |
|
24 |
-
@spaces.GPU
|
25 |
def load_model():
|
26 |
-
"""Load MiniCPM model
|
27 |
global _model, _tokenizer
|
28 |
|
29 |
if _model is not None and _tokenizer is not None:
|
@@ -39,7 +38,7 @@ def load_model():
|
|
39 |
"openbmb/MiniCPM-V-2_6",
|
40 |
trust_remote_code=True,
|
41 |
torch_dtype=torch.float16,
|
42 |
-
device_map="auto"
|
43 |
)
|
44 |
return _model, _tokenizer
|
45 |
except Exception as e:
|
@@ -58,7 +57,7 @@ def load_model():
|
|
58 |
return _model, _tokenizer
|
59 |
|
60 |
def pdf_to_images(pdf_file):
|
61 |
-
"""Convert PDF file to list of PIL images"""
|
62 |
try:
|
63 |
if hasattr(pdf_file, 'read'):
|
64 |
pdf_bytes = pdf_file.read()
|
@@ -151,12 +150,9 @@ INSTRUCTIONS:
|
|
151 |
6. If information is not visible, leave field empty but still include it
|
152 |
7. Return ONLY the JSON, no other text"""
|
153 |
|
154 |
-
|
155 |
-
|
156 |
-
"""Extract data from a single image using MiniCPM on GPU"""
|
157 |
try:
|
158 |
-
model, tokenizer = load_model()
|
159 |
-
|
160 |
# Convert PIL image to proper format if needed
|
161 |
if hasattr(image, 'convert'):
|
162 |
image = image.convert('RGB')
|
@@ -169,7 +165,7 @@ def extract_data_from_image(image, extraction_prompt):
|
|
169 |
"content": extraction_prompt
|
170 |
}],
|
171 |
tokenizer=tokenizer,
|
172 |
-
sampling=False,
|
173 |
temperature=0.1,
|
174 |
max_new_tokens=2048
|
175 |
)
|
@@ -274,9 +270,9 @@ def combine_page_data(pages_data):
|
|
274 |
}
|
275 |
}
|
276 |
|
277 |
-
@spaces.GPU(duration=
|
278 |
def extract_efax_from_pdf(pdf_file, custom_prompt=None):
|
279 |
-
"""Main function to process multi-page PDF eFax
|
280 |
try:
|
281 |
if pdf_file is None:
|
282 |
return {
|
@@ -286,7 +282,8 @@ def extract_efax_from_pdf(pdf_file, custom_prompt=None):
|
|
286 |
"pages_data": []
|
287 |
}
|
288 |
|
289 |
-
# Convert PDF to images
|
|
|
290 |
images = pdf_to_images(pdf_file)
|
291 |
|
292 |
if not images:
|
@@ -297,30 +294,38 @@ def extract_efax_from_pdf(pdf_file, custom_prompt=None):
|
|
297 |
"pages_data": []
|
298 |
}
|
299 |
|
300 |
-
|
|
|
|
|
|
|
|
|
|
|
301 |
extraction_prompt = custom_prompt if custom_prompt else get_medical_extraction_prompt()
|
302 |
|
303 |
-
# Process
|
304 |
pages_data = []
|
305 |
for i, image in enumerate(images):
|
306 |
-
print(f"Processing page {i+1}/{len(images)}")
|
307 |
-
page_result = extract_data_from_image(image, extraction_prompt)
|
308 |
pages_data.append({
|
309 |
"page_number": i + 1,
|
310 |
"page_data": page_result
|
311 |
})
|
312 |
|
313 |
-
|
|
|
|
|
314 |
combined_result = combine_page_data(pages_data)
|
315 |
|
316 |
-
# Final result
|
317 |
result = {
|
318 |
"status": "success",
|
319 |
"total_pages": len(images),
|
320 |
"pages_data": pages_data,
|
321 |
"combined_extraction": combined_result,
|
322 |
"model_used": "MiniCPM-V-2_6-ZeroGPU",
|
323 |
-
"hardware": "ZeroGPU"
|
|
|
324 |
}
|
325 |
|
326 |
return result
|
@@ -335,16 +340,16 @@ def extract_efax_from_pdf(pdf_file, custom_prompt=None):
|
|
335 |
|
336 |
# Create Gradio Interface
|
337 |
def create_gradio_interface():
|
338 |
-
with gr.Blocks(title="eFax PDF Data Extractor - ZeroGPU", theme=gr.themes.Soft()) as demo:
|
339 |
gr.Markdown("# π₯ eFax Medical Data Extraction API")
|
340 |
-
gr.Markdown("π **GPU
|
341 |
|
342 |
with gr.Tab("π PDF Upload & Extraction"):
|
343 |
with gr.Row():
|
344 |
with gr.Column():
|
345 |
pdf_input = gr.File(
|
346 |
file_types=[".pdf"],
|
347 |
-
label="Upload eFax PDF",
|
348 |
file_count="single"
|
349 |
)
|
350 |
|
@@ -356,7 +361,14 @@ def create_gradio_interface():
|
|
356 |
placeholder="Leave empty to use optimized medical data extraction prompt..."
|
357 |
)
|
358 |
|
359 |
-
extract_btn = gr.Button("π Extract Medical Data (GPU)", variant="primary", size="lg")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
360 |
|
361 |
with gr.Column():
|
362 |
status_output = gr.Textbox(label="π Processing Status", interactive=False)
|
@@ -364,17 +376,14 @@ def create_gradio_interface():
|
|
364 |
|
365 |
with gr.Tab("π API Usage"):
|
366 |
gr.Markdown("""
|
367 |
-
## API
|
368 |
-
|
369 |
-
Your Space runs on **ZeroGPU** for 10-50x faster processing!
|
370 |
|
371 |
-
### Python
|
372 |
```
|
373 |
import requests
|
374 |
import base64
|
375 |
|
376 |
-
|
377 |
-
with open("medical_fax.pdf", "rb") as f:
|
378 |
pdf_b64 = base64.b64encode(f.read()).decode()
|
379 |
|
380 |
response = requests.post(
|
@@ -382,62 +391,39 @@ def create_gradio_interface():
|
|
382 |
json={
|
383 |
"data": [
|
384 |
{"name": "medical_fax.pdf", "data": f"application/pdf;base64,{pdf_b64}"},
|
385 |
-
"" #
|
386 |
]
|
387 |
}
|
388 |
)
|
389 |
|
|
|
390 |
result = response.json()
|
391 |
-
|
392 |
-
# Access combined medical data
|
393 |
medical_data = result["data"]["combined_extraction"]
|
394 |
-
print("Patient:", medical_data["data"]["patient_first_name"], medical_data["data"]["patient_last_name"])
|
395 |
-
print("Insurance:", medical_data["data"]["primary_insurance"]["payer_name"])
|
396 |
-
```
|
397 |
-
|
398 |
-
### Response Format
|
399 |
-
```
|
400 |
-
{
|
401 |
-
"status": "success",
|
402 |
-
"total_pages": 13,
|
403 |
-
"combined_extraction": {
|
404 |
-
"data": {
|
405 |
-
"patient_first_name": "John",
|
406 |
-
"patient_last_name": "Doe",
|
407 |
-
"primary_insurance": {
|
408 |
-
"payer_name": "UNITED HEALTHCARE",
|
409 |
-
"member_id": "123456789"
|
410 |
-
}
|
411 |
-
},
|
412 |
-
"confidence_scores": {...},
|
413 |
-
"metadata": {...}
|
414 |
-
}
|
415 |
-
}
|
416 |
```
|
417 |
""")
|
418 |
|
419 |
with gr.Tab("β‘ Performance Info"):
|
420 |
gr.Markdown("""
|
421 |
-
## ZeroGPU Performance
|
422 |
|
423 |
-
|
424 |
-
-
|
425 |
-
-
|
426 |
-
-
|
427 |
-
- **π‘ Dynamic Allocation**: GPU activates only during processing
|
428 |
|
429 |
-
|
430 |
-
-
|
431 |
-
-
|
432 |
-
-
|
433 |
-
-
|
434 |
-
- β
Referral Source & Priority
|
435 |
-
- β
Confidence Scores for Quality Control
|
436 |
|
437 |
-
|
438 |
-
|
439 |
-
|
440 |
-
|
|
|
|
|
|
|
441 |
""")
|
442 |
|
443 |
def process_with_status(pdf_file, custom_prompt):
|
@@ -450,7 +436,7 @@ def create_gradio_interface():
|
|
450 |
result = extract_efax_from_pdf(pdf_file, custom_prompt if custom_prompt.strip() else None)
|
451 |
|
452 |
if result["status"] == "success":
|
453 |
-
yield f"β
Successfully processed {result['total_pages']} pages", result
|
454 |
else:
|
455 |
yield f"β Error: {result.get('error', 'Unknown error')}", result
|
456 |
|
@@ -477,4 +463,4 @@ if __name__ == "__main__":
|
|
477 |
server_name="0.0.0.0",
|
478 |
server_port=7860,
|
479 |
show_error=True
|
480 |
-
)
|
|
|
21 |
_model = None
|
22 |
_tokenizer = None
|
23 |
|
|
|
24 |
def load_model():
|
25 |
+
"""Load MiniCPM model (CPU loading, GPU usage happens in main function)"""
|
26 |
global _model, _tokenizer
|
27 |
|
28 |
if _model is not None and _tokenizer is not None:
|
|
|
38 |
"openbmb/MiniCPM-V-2_6",
|
39 |
trust_remote_code=True,
|
40 |
torch_dtype=torch.float16,
|
41 |
+
device_map="auto" # Will move to GPU when @spaces.GPU is active
|
42 |
)
|
43 |
return _model, _tokenizer
|
44 |
except Exception as e:
|
|
|
57 |
return _model, _tokenizer
|
58 |
|
59 |
def pdf_to_images(pdf_file):
|
60 |
+
"""Convert PDF file to list of PIL images (CPU operation)"""
|
61 |
try:
|
62 |
if hasattr(pdf_file, 'read'):
|
63 |
pdf_bytes = pdf_file.read()
|
|
|
150 |
6. If information is not visible, leave field empty but still include it
|
151 |
7. Return ONLY the JSON, no other text"""
|
152 |
|
153 |
+
def extract_data_from_image(image, extraction_prompt, model, tokenizer):
|
154 |
+
"""Extract data from a single image using MiniCPM (runs within GPU session)"""
|
|
|
155 |
try:
|
|
|
|
|
156 |
# Convert PIL image to proper format if needed
|
157 |
if hasattr(image, 'convert'):
|
158 |
image = image.convert('RGB')
|
|
|
165 |
"content": extraction_prompt
|
166 |
}],
|
167 |
tokenizer=tokenizer,
|
168 |
+
sampling=False,
|
169 |
temperature=0.1,
|
170 |
max_new_tokens=2048
|
171 |
)
|
|
|
270 |
}
|
271 |
}
|
272 |
|
273 |
+
@spaces.GPU(duration=600) # 10 minutes for large documents
|
274 |
def extract_efax_from_pdf(pdf_file, custom_prompt=None):
|
275 |
+
"""Main function to process multi-page PDF eFax - ALL GPU processing happens here"""
|
276 |
try:
|
277 |
if pdf_file is None:
|
278 |
return {
|
|
|
282 |
"pages_data": []
|
283 |
}
|
284 |
|
285 |
+
# Step 1: Convert PDF to images (CPU operation - do this before GPU)
|
286 |
+
print("Converting PDF to images...")
|
287 |
images = pdf_to_images(pdf_file)
|
288 |
|
289 |
if not images:
|
|
|
294 |
"pages_data": []
|
295 |
}
|
296 |
|
297 |
+
print(f"Converted {len(images)} pages. Starting GPU processing...")
|
298 |
+
|
299 |
+
# Step 2: Load model on GPU (happens once GPU session starts)
|
300 |
+
model, tokenizer = load_model()
|
301 |
+
|
302 |
+
# Step 3: Use custom prompt or default
|
303 |
extraction_prompt = custom_prompt if custom_prompt else get_medical_extraction_prompt()
|
304 |
|
305 |
+
# Step 4: Process all pages within single GPU session
|
306 |
pages_data = []
|
307 |
for i, image in enumerate(images):
|
308 |
+
print(f"Processing page {i+1}/{len(images)} on GPU...")
|
309 |
+
page_result = extract_data_from_image(image, extraction_prompt, model, tokenizer)
|
310 |
pages_data.append({
|
311 |
"page_number": i + 1,
|
312 |
"page_data": page_result
|
313 |
})
|
314 |
|
315 |
+
print("GPU processing complete. Combining results...")
|
316 |
+
|
317 |
+
# Step 5: Combine data from all pages
|
318 |
combined_result = combine_page_data(pages_data)
|
319 |
|
320 |
+
# Final result
|
321 |
result = {
|
322 |
"status": "success",
|
323 |
"total_pages": len(images),
|
324 |
"pages_data": pages_data,
|
325 |
"combined_extraction": combined_result,
|
326 |
"model_used": "MiniCPM-V-2_6-ZeroGPU",
|
327 |
+
"hardware": "ZeroGPU",
|
328 |
+
"processing_time": "Within 10-minute GPU session"
|
329 |
}
|
330 |
|
331 |
return result
|
|
|
340 |
|
341 |
# Create Gradio Interface
|
342 |
def create_gradio_interface():
|
343 |
+
with gr.Blocks(title="eFax PDF Data Extractor - Optimized ZeroGPU", theme=gr.themes.Soft()) as demo:
|
344 |
gr.Markdown("# π₯ eFax Medical Data Extraction API")
|
345 |
+
gr.Markdown("π **Optimized GPU Usage** - Single 10-minute GPU session for entire document")
|
346 |
|
347 |
with gr.Tab("π PDF Upload & Extraction"):
|
348 |
with gr.Row():
|
349 |
with gr.Column():
|
350 |
pdf_input = gr.File(
|
351 |
file_types=[".pdf"],
|
352 |
+
label="Upload eFax PDF (up to 20 pages)",
|
353 |
file_count="single"
|
354 |
)
|
355 |
|
|
|
361 |
placeholder="Leave empty to use optimized medical data extraction prompt..."
|
362 |
)
|
363 |
|
364 |
+
extract_btn = gr.Button("π Extract Medical Data (10min GPU)", variant="primary", size="lg")
|
365 |
+
|
366 |
+
gr.Markdown("""
|
367 |
+
### β‘ Optimized Processing
|
368 |
+
- **Single GPU Session**: All pages processed in one 10-minute session
|
369 |
+
- **No Timeouts**: Handles up to 20+ page documents
|
370 |
+
- **Efficient**: PDFβImages (CPU) β All Processing (GPU) β Results
|
371 |
+
""")
|
372 |
|
373 |
with gr.Column():
|
374 |
status_output = gr.Textbox(label="π Processing Status", interactive=False)
|
|
|
376 |
|
377 |
with gr.Tab("π API Usage"):
|
378 |
gr.Markdown("""
|
379 |
+
## Optimized API (No Timeout Issues)
|
|
|
|
|
380 |
|
381 |
+
### Python Usage
|
382 |
```
|
383 |
import requests
|
384 |
import base64
|
385 |
|
386 |
+
with open("large_medical_fax.pdf", "rb") as f:
|
|
|
387 |
pdf_b64 = base64.b64encode(f.read()).decode()
|
388 |
|
389 |
response = requests.post(
|
|
|
391 |
json={
|
392 |
"data": [
|
393 |
{"name": "medical_fax.pdf", "data": f"application/pdf;base64,{pdf_b64}"},
|
394 |
+
"" # Empty for default prompt
|
395 |
]
|
396 |
}
|
397 |
)
|
398 |
|
399 |
+
# Now handles 13+ pages without timeout!
|
400 |
result = response.json()
|
|
|
|
|
401 |
medical_data = result["data"]["combined_extraction"]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
402 |
```
|
403 |
""")
|
404 |
|
405 |
with gr.Tab("β‘ Performance Info"):
|
406 |
gr.Markdown("""
|
407 |
+
## Optimized ZeroGPU Performance
|
408 |
|
409 |
+
### Before Optimization (β Had Timeout Issues)
|
410 |
+
- GPU session per page = 13 Γ 30 seconds = 6.5 minutes
|
411 |
+
- Model loading repeated = wasted time
|
412 |
+
- Timeout around page 11/13
|
|
|
413 |
|
414 |
+
### After Optimization (β
No Timeouts)
|
415 |
+
- **Single 10-minute GPU session** for entire document
|
416 |
+
- Model loads once, processes all pages
|
417 |
+
- Handles 15-20+ page documents easily
|
418 |
+
- PDF conversion on CPU (doesn't count toward GPU time)
|
|
|
|
|
419 |
|
420 |
+
### Processing Flow
|
421 |
+
1. **PDF β Images** (CPU, before GPU starts)
|
422 |
+
2. **π GPU Session Starts** (10 minutes allocated)
|
423 |
+
3. **Load Model** (once, on GPU)
|
424 |
+
4. **Process All Pages** (GPU, sequential)
|
425 |
+
5. **GPU Session Ends**
|
426 |
+
6. **Combine Results** (CPU, after GPU)
|
427 |
""")
|
428 |
|
429 |
def process_with_status(pdf_file, custom_prompt):
|
|
|
436 |
result = extract_efax_from_pdf(pdf_file, custom_prompt if custom_prompt.strip() else None)
|
437 |
|
438 |
if result["status"] == "success":
|
439 |
+
yield f"β
Successfully processed {result['total_pages']} pages in single GPU session", result
|
440 |
else:
|
441 |
yield f"β Error: {result.get('error', 'Unknown error')}", result
|
442 |
|
|
|
463 |
server_name="0.0.0.0",
|
464 |
server_port=7860,
|
465 |
show_error=True
|
466 |
+
)
|