import gradio as gr
from transformers import pipeline

# Load the Microsoft BioGPT text generation pipeline
generator = pipeline("text-generation", model="microsoft/BioGPT")

def generate_medical_text(prompt, temperature, top_p, max_length):
    # Clamp temperature to at least 0.01 to avoid division by zero errors
    temperature = max(0.01, float(temperature))
    
    # Generate text using the pipeline
    results = generator(
        prompt,
        max_length=int(max_length),
        temperature=temperature,
        top_p=float(top_p),
        do_sample=True,
        num_return_sequences=1,
        truncation=True,
        pad_token_id=generator.tokenizer.eos_token_id # Prevents padding warnings
    )
    
    return results[0]["generated_text"]

# Define interface text
title = "Medical Text Generator"
description = (
    "Designed for experimenting with medical and health-related text — "
    "clinical notes, symptom descriptions, patient scenarios, and health explanations. "
    "Powered by Microsoft's `BioGPT` model, which was trained on millions of biomedical research articles."
)

# Define preset examples (Format: [prompt, temperature, top_p, max_length])
examples = [["The patient presented with symptoms of", 0.5, 0.9, 120],["Common side effects of this medication include", 0.5, 0.9, 120],["The doctor examined the test results and concluded", 0.5, 0.9, 120],["A healthy diet for someone with diabetes should", 0.5, 0.9, 120],["The difference between a virus and a bacteria is", 0.5, 0.9, 120]
]

# Build the Gradio interface
demo = gr.Interface(
    fn=generate_medical_text,
    inputs=[
        gr.Textbox(
            lines=3, 
            label="Prompt", 
            placeholder="Enter a medical prompt here..."
        ),
        gr.Slider(
            minimum=0.1, maximum=2.0, value=0.5, step=0.1, 
            label="Temperature", 
            info="Controls how creative/wild the writing is"
        ),
        gr.Slider(
            minimum=0.1, maximum=1.0, value=0.9, step=0.05, 
            label="Top-p", 
            info="Controls word diversity"
        ),
        gr.Slider(
            minimum=20, maximum=200, value=120, step=1, 
            label="Max Length", 
            info="Controls how much text it generates"
        )
    ],
    outputs=gr.Textbox(label="Generated Text", lines=8),
    title=title,
    description=description,
    examples=examples
)

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