Spaces:
Paused
Paused
| #!/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, | |
| ) | |
| from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| if device == "cuda": | |
| attn_implementation = "sdpa" | |
| dtype = torch.bfloat16 | |
| print("Using sdpa for GPU") | |
| else: | |
| attn_implementation = "eager" | |
| dtype = torch.float32 | |
| print("Using eager attention for CPU") | |
| print(f"Loading LightOnOCR model on {device} with {attn_implementation} attention...") | |
| ocr_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("LightOnOCR model loaded successfully!") | |
| # -------- Clinical NER models (load ONCE) -------- | |
| print("Loading clinical NER model...") | |
| ner_tokenizer = AutoTokenizer.from_pretrained("samrawal/bert-base-uncased_clinical-ner") | |
| ner_model = AutoModelForTokenClassification.from_pretrained("samrawal/bert-base-uncased_clinical-ner") | |
| ner_pipeline = pipeline( | |
| "ner", | |
| model=ner_model, | |
| tokenizer=ner_tokenizer, | |
| aggregation_strategy="simple", | |
| ) | |
| print("Clinical NER model loaded successfully!") | |
| 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.""" | |
| # Remove common chat template markers | |
| markers_to_remove = ["system", "user", "assistant"] | |
| # Split by lines and filter | |
| 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) | |
| # Join back and strip leading/trailing whitespace | |
| cleaned = '\n'.join(cleaned_lines).strip() | |
| # Alternative approach: if there's an "assistant" marker, take everything after it | |
| if "assistant" in text.lower(): | |
| parts = text.split("assistant", 1) | |
| if len(parts) > 1: | |
| cleaned = parts[1].strip() | |
| return cleaned | |
| def extract_text_from_image(image, temperature=0.2, stream=False): | |
| """Extract text from image using LightOnOCR model, and run clinical NER.""" | |
| # Prepare the chat format | |
| chat = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "image", "url": image}, # adjust to {"type": "image", "image": image} if LightOnOCR expects that | |
| ], | |
| } | |
| ] | |
| # Tokenize | |
| inputs = processor.apply_chat_template( | |
| chat, | |
| add_generation_prompt=True, | |
| tokenize=True, | |
| return_dict=True, | |
| return_tensors="pt", | |
| ) | |
| # Move inputs to device | |
| 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=ocr_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) | |
| # For streaming, we’ll only show text progressively, | |
| # and keep medications empty (or compute at the end if you prefer). | |
| yield cleaned_text, "" | |
| thread.join() | |
| else: | |
| # Non-streaming generation | |
| with torch.no_grad(): | |
| outputs = ocr_model.generate(**generation_kwargs) | |
| output_text = processor.decode(outputs[0], skip_special_tokens=True) | |
| cleaned_text = clean_output_text(output_text) | |
| # Clinical NER on the full cleaned text | |
| entities = ner_pipeline(cleaned_text) | |
| medications = [] | |
| for ent in entities: | |
| if ent["entity_group"] == "treatment": | |
| word = ent["word"] | |
| if word.startswith("##") and medications: | |
| medications[-1] += word[2:] | |
| else: | |
| medications.append(word) | |
| medications_str = ", ".join(set(medications)) if medications else "None detected" | |
| yield cleaned_text, medications_str | |
| def process_input(file_input, temperature, page_num, enable_streaming): | |
| """Process uploaded file (image or PDF) and extract text with optional streaming.""" | |
| if file_input is None: | |
| # 6 outputs: [output_text, medications_output, raw_output, page_info, rendered_image, num_pages] | |
| yield "Please upload an image or PDF first.", "", "", "", None, 1 | |
| return | |
| image_to_process = None | |
| page_info = "" | |
| slider_value = page_num | |
| 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}" | |
| slider_value = actual_page | |
| except Exception as e: | |
| msg = f"Error processing PDF: {str(e)}" | |
| yield msg, "", msg, "", None, slider_value | |
| return | |
| else: | |
| # Handle image files | |
| try: | |
| image_to_process = Image.open(file_path) | |
| page_info = "Processing image" | |
| except Exception as e: | |
| msg = f"Error opening image: {str(e)}" | |
| yield msg, "", msg, "", None, slider_value | |
| return | |
| try: | |
| # Extract text using LightOnOCR with optional streaming | |
| for extracted_text, medications in extract_text_from_image( | |
| image_to_process, temperature, stream=enable_streaming | |
| ): | |
| raw_md = extracted_text # or you can keep a different raw version | |
| # 6 outputs: markdown_text, medications, raw_output, page_info, image, slider | |
| yield extracted_text, medications, raw_md, page_info, image_to_process, gr.update( | |
| value=slider_value | |
| ) | |
| except Exception as e: | |
| error_msg = f"Error during text extraction: {str(e)}" | |
| # 6 outputs | |
| yield error_msg, "", error_msg, page_info, image_to_process, gr.update(value=slider_value) | |
| 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) | |
| # Create Gradio interface | |
| with gr.Blocks(title="📖 Image/PDF OCR with LightOnOCR", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown(f""" | |
| # 📖 Image/PDF to Text Extraction with LightOnOCR | |
| **💡 How to use:** | |
| 1. Upload an image or PDF | |
| 2. For PDFs: select which page to extract (1-20) | |
| 3. Adjust temperature if needed | |
| 4. Click "Extract Text" | |
| **Note:** The Markdown rendering for tables may not always be perfect. Check the raw output for complex tables! | |
| **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 Text", variant="primary") | |
| clear_btn = gr.Button("Clear", variant="secondary") | |
| with gr.Column(scale=2): | |
| output_text = gr.Markdown( | |
| label="📄 Extracted Text (Rendered)", | |
| value="*Extracted text will appear here...*" | |
| ) | |
| medications_output = gr.Textbox( | |
| label="💊 Extracted Medicines/Drugs", | |
| placeholder="Medicine/drug names will appear here...", | |
| lines=2, | |
| max_lines=5, | |
| interactive=False, | |
| show_copy_button=True | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| raw_output = gr.Textbox( | |
| label="Raw Markdown Output", | |
| placeholder="Raw text 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, medications_output, 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, "*Extracted text will appear here...*", "", "", None, 1), | |
| outputs=[file_input, output_text, raw_output, page_info, rendered_image, num_pages] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |
| #################################### old code to be checked ############################################# | |
| # 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() | |