app.py

#!/usr/bin/env python3
import subprocess
import sys
import threading

import spaces
import torch

import gradio as gr
from PIL import Image
from io import BytesIO
import pypdfium2 as pdfium
from transformers import (
    LightOnOCRForConditionalGeneration,
    LightOnOCRProcessor,
    TextIteratorStreamer,
)

# ---- CLINICAL NER IMPORTS ----
import spacy

device = "cuda" if torch.cuda.is_available() else "cpu"

# Choose best attention implementation based on device
if device == "cuda":
    attn_implementation = "sdpa"  
    dtype = torch.bfloat16
    print("Using sdpa for GPU")
else:
    attn_implementation = "eager"  # Best for CPU
    dtype = torch.float32
    print("Using eager attention for CPU")

# Initialize the LightOnOCR model and processor
print(f"Loading model on {device} with {attn_implementation} attention...")
model = LightOnOCRForConditionalGeneration.from_pretrained(
    "lightonai/LightOnOCR-1B-1025",
    attn_implementation=attn_implementation,
    torch_dtype=dtype,
    trust_remote_code=True
).to(device).eval()

processor = LightOnOCRProcessor.from_pretrained(
    "lightonai/LightOnOCR-1B-1025",
    trust_remote_code=True
)
print("Model loaded successfully!")

# ---- LOAD CLINICAL NER MODEL (BC5CDR) ----
print("Loading clinical NER model (bc5cdr)...")
nlp_ner = spacy.load("en_ner_bc5cdr_md")
print("Clinical NER loaded.")

def render_pdf_page(page, max_resolution=1540, scale=2.77):
    """Render a PDF page to PIL Image."""
    width, height = page.get_size()
    pixel_width = width * scale
    pixel_height = height * scale
    resize_factor = min(1, max_resolution / pixel_width, max_resolution / pixel_height)
    target_scale = scale * resize_factor
    return page.render(scale=target_scale, rev_byteorder=True).to_pil()


def process_pdf(pdf_path, page_num=1):
    """Extract a specific page from PDF."""
    pdf = pdfium.PdfDocument(pdf_path)
    total_pages = len(pdf)
    page_idx = min(max(int(page_num) - 1, 0), total_pages - 1)
    
    page = pdf[page_idx]
    img = render_pdf_page(page)
    
    pdf.close()
    return img, total_pages, page_idx + 1


def clean_output_text(text):
    """Remove chat template artifacts from output."""
    markers_to_remove = ["system", "user", "assistant"]
    lines = text.split('\n')
    cleaned_lines = []
    for line in lines:
        stripped = line.strip()
        # Skip lines that are just template markers
        if stripped.lower() not in markers_to_remove:
            cleaned_lines.append(line)
    cleaned = '\n'.join(cleaned_lines).strip()
    if "assistant" in text.lower():
        parts = text.split("assistant", 1)
        if len(parts) > 1:
            cleaned = parts[1].strip()
    return cleaned

def extract_medication_names(text):
    """Extract medication names using clinical NER (spacy: bc5cdr CHEMICAL)."""
    doc = nlp_ner(text)
    meds = [ent.text for ent in doc.ents if ent.label_ == "CHEMICAL"]
    meds_unique = list(dict.fromkeys(meds))
    return meds_unique


@spaces.GPU
def extract_text_from_image(image, temperature=0.2, stream=False):
    """Extract text from image using LightOnOCR model."""
    chat = [
        {
            "role": "user",
            "content": [
                {"type": "image", "url": image},
            ],
        }
    ]
    inputs = processor.apply_chat_template(
        chat,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt"
    )
    inputs = {
        k: v.to(device=device, dtype=dtype) if isinstance(v, torch.Tensor) and v.dtype in [torch.float32, torch.float16, torch.bfloat16]
        else v.to(device) if isinstance(v, torch.Tensor) 
        else v 
        for k, v in inputs.items()
    }
    generation_kwargs = dict(
        **inputs,
        max_new_tokens=2048,
        temperature=temperature if temperature > 0 else 0.0,
        use_cache=True,
        do_sample=temperature > 0,
    )
    if stream:
        # Streaming generation
        streamer = TextIteratorStreamer(
            processor.tokenizer,
            skip_prompt=True,
            skip_special_tokens=True
        )
        generation_kwargs["streamer"] = streamer
        thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
        thread.start()
        full_text = ""
        for new_text in streamer:
            full_text += new_text
            cleaned_text = clean_output_text(full_text)
            yield cleaned_text
        thread.join()
    else:
        # Non-streaming generation
        with torch.no_grad():
            outputs = model.generate(**generation_kwargs)
        output_text = processor.decode(outputs[0], skip_special_tokens=True)
        cleaned_text = clean_output_text(output_text)
        yield cleaned_text

def process_input(file_input, temperature, page_num, enable_streaming):
    """Process uploaded file (image or PDF) and extract medication names via OCR+NER."""
    if file_input is None:
        yield "Please upload an image or PDF first.", "", "", None, gr.update()
        return
    image_to_process = None
    page_info = ""
    file_path = file_input if isinstance(file_input, str) else file_input.name
    # Handle PDF files
    if file_path.lower().endswith('.pdf'):
        try:
            image_to_process, total_pages, actual_page = process_pdf(file_path, int(page_num))
            page_info = f"Processing page {actual_page} of {total_pages}"
        except Exception as e:
            yield f"Error processing PDF: {str(e)}", "", "", None, gr.update()
            return
    # Handle image files
    else:
        try:
            image_to_process = Image.open(file_path)
            page_info = "Processing image"
        except Exception as e:
            yield f"Error opening image: {str(e)}", "", "", None, gr.update()
            return
    try:
        for extracted_text in extract_text_from_image(image_to_process, temperature, stream=enable_streaming):
            meds = extract_medication_names(extracted_text)
            meds_str = "\n".join(meds) if meds else "No medications found."
            yield meds_str, meds_str, page_info, image_to_process, gr.update()
    except Exception as e:
        error_msg = f"Error during text extraction: {str(e)}"
        yield error_msg, error_msg, page_info, image_to_process, gr.update()

def update_slider(file_input):
    """Update page slider based on PDF page count."""
    if file_input is None:
        return gr.update(maximum=20, value=1)
    file_path = file_input if isinstance(file_input, str) else file_input.name
    if file_path.lower().endswith('.pdf'):
        try:
            pdf = pdfium.PdfDocument(file_path)
            total_pages = len(pdf)
            pdf.close()
            return gr.update(maximum=total_pages, value=1)
        except:
            return gr.update(maximum=20, value=1)
    else:
        return gr.update(maximum=1, value=1)

# ----- GRADIO UI -----
with gr.Blocks(title="📖 Image/PDF OCR + Clinical NER", theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"""
# 📖 Medication Extraction from Image/PDF with LightOnOCR + Clinical NER

**💡 How to use:**
1. Upload an image or PDF
2. For PDFs: select which page to extract
3. Adjust temperature if needed
4. Click "Extract Medications"

**Output:** Only medication names found in text (via NER)

**Model:** LightOnOCR-1B-1025 by LightOn AI  
**Device:** {device.upper()}  
**Attention:** {attn_implementation}
""")
    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(
                label="🖼️ Upload Image or PDF",
                file_types=[".pdf", ".png", ".jpg", ".jpeg"],
                type="filepath"
            )
            rendered_image = gr.Image(
                label="📄 Preview",
                type="pil",
                height=400,
                interactive=False
            )
            num_pages = gr.Slider(
                minimum=1,
                maximum=20,
                value=1,
                step=1,
                label="PDF: Page Number",
                info="Select which page to extract"
            )
            page_info = gr.Textbox(
                label="Processing Info",
                value="",
                interactive=False
            )
            temperature = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.2,
                step=0.05,
                label="Temperature",
                info="0.0 = deterministic, Higher = more varied"
            )
            enable_streaming = gr.Checkbox(
                label="Enable Streaming",
                value=True,
                info="Show text progressively as it's generated"
            )
            submit_btn = gr.Button("Extract Medications", variant="primary")
            clear_btn = gr.Button("Clear", variant="secondary")
        with gr.Column(scale=2):
            output_text = gr.Markdown(
                label="🩺 Extracted Medication Names",
                value="*Medication names will appear here...*"
            )
    with gr.Row():
        with gr.Column():
            raw_output = gr.Textbox(
                label="Extracted Medication Names (Raw)",
                placeholder="Medication list will appear here...",
                lines=20,
                max_lines=30,
                show_copy_button=True
            )
    # Event handlers
    submit_btn.click(
        fn=process_input,
        inputs=[file_input, temperature, num_pages, enable_streaming],
        outputs=[output_text, raw_output, page_info, rendered_image, num_pages]
    )
    file_input.change(
        fn=update_slider,
        inputs=[file_input],
        outputs=[num_pages]
    )
    clear_btn.click(
        fn=lambda: (None, "*Medication names will appear here...*", "", "", None, 1),
        outputs=[file_input, output_text, raw_output, page_info, rendered_image, num_pages]
    )

if __name__ == "__main__":
    demo.launch()