import gradio as gr
from transformers import AutoTokenizer, AutoModelForMaskedLM, pipeline

# Load BlueBERT model and tokenizer
model_name = "bionlp/bluebert_pubmed_mimic_uncased_L-12_H-768_A-12"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForMaskedLM.from_pretrained(model_name)

# Create a fill-mask pipeline to handle predictions
nlp = pipeline("fill-mask", model=model, tokenizer=tokenizer)

def predict(text):
    # Check if the input text contains the [MASK] token
    if "[MASK]" not in text:
        return "Error: Please enter a sentence containing a [MASK] token."
    
    # Process the text using the model
    result = nlp(text)
    return result

# Gradio interface to input text and output model predictions
iface = gr.Interface(
    fn=predict, 
    inputs=gr.Textbox(lines=2, placeholder="Enter a sentence with [MASK]..."), 
    outputs="json",
    title="BlueBERT Testing",
    description="Test BlueBERT on biomedical data or general text"
)

iface.launch()