ai_med_extract/gradio_app.py

import gradio as gr
import os
from .app import agents
import tempfile
import logging

logging.basicConfig(level=logging.INFO)

def process_document(file, process_type):
    try:
        # Create a temporary file to store the upload
        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.name)[1]) as temp_file:
            temp_file.write(file.read())
            temp_path = temp_file.name

        results = {}
        
        if process_type == "text_extraction":
            results["extracted_text"] = agents["text_extractor"].extract_text(temp_path)
            results["phi_scrubbed"] = agents["phi_scrubber"].scrub_phi(results["extracted_text"])
            
        elif process_type == "medical_data":
            text = agents["text_extractor"].extract_text(temp_path)
            results["medical_data"] = agents["medical_data_extractor"].extract_medical_data(text)
            
        elif process_type == "summarization":
            text = agents["text_extractor"].extract_text(temp_path)
            results["summary"] = agents["summarizer"].summarize(text)
            
        elif process_type == "audio_transcription":
            results["transcription"] = agents["whisper_model"].transcribe(temp_path)
            
        # Clean up temporary file
        os.unlink(temp_path)
        
        return results
        
    except Exception as e:
        logging.error(f"Error processing document: {str(e)}", exc_info=True)
        return {"error": str(e)}

# Create the Gradio interface
def create_interface():
    with gr.Blocks(title="Medical Document Processor") as interface:
        gr.Markdown("# Medical Document Processor")
        gr.Markdown("Upload your medical document and select the processing type.")
        
        with gr.Row():
            with gr.Column():
                file_input = gr.File(label="Upload Document")
                process_type = gr.Radio(
                    choices=["text_extraction", "medical_data", "summarization", "audio_transcription"],
                    label="Processing Type"
                )
                process_btn = gr.Button("Process Document")
            
            with gr.Column():
                output = gr.JSON(label="Results")
        
        process_btn.click(
            fn=process_document,
            inputs=[file_input, process_type],
            outputs=output
        )
    
    return interface

# Create and launch the interface
interface = create_interface()
interface.launch(server_name="0.0.0.0", server_port=7860)