import gradio as gr
import re

def extract_study_timeline(text):
    """Extracts Screening, Treatment, and Follow-up durations from study timeline text."""
    screening = re.search(r'Screening\s*(\d+)\s*weeks?', text)
    treatment = re.search(r'Treatment\s*(\d+)\s*weeks?', text)
    follow_up = re.search(r'Follow[-\s]*up\s*(\d+)\s*weeks?', text)

    timeline = {
        "Screening": int(screening.group(1)) if screening else None,
        "Treatment": int(treatment.group(1)) if treatment else None,
        "Follow-Up": int(follow_up.group(1)) if follow_up else None
    }

    return timeline

# Create a Gradio Interface
demo = gr.Interface(
    fn=extract_study_timeline,
    inputs=gr.Textbox(lines=3, placeholder="Enter study timeline text..."),
    outputs="json",
    title="Study Timeline Extraction",
    description="Enter a study timeline description, and the model will extract Screening, Treatment, and Follow-Up durations.",
)

# Launch the Gradio App
demo.launch()