from transformers import pipeline, set_seed
from transformers import BioGptTokenizer, BioGptForCausalLM
import gradio as gr

model_list = [
    "microsoft/biogpt",
    "microsoft/BioGPT-Large-PubMedQA"
]

def biogpt(
    prompt: str,
    model_id: str,
    max_length: int = 25,
    num_return_sequences: int = 5
):
    
    model = BioGptForCausalLM.from_pretrained(model_id)
    tokenizer = BioGptTokenizer.from_pretrained(model_id)
    generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
    set_seed(42)
    output = generator(prompt, max_length=max_length, num_return_sequences=num_return_sequences, do_sample=True)
    output_dict = {
        "1": output[0]['generated_text'],
        "2": output[1]['generated_text'],
        "3": output[2]['generated_text'],
        "4": output[3]['generated_text'],
        "5": output[4]['generated_text']
    }
    return f'{output_dict["1"]}\n\n{output_dict["2"]}\n\n{output_dict["3"]}\n\n{output_dict["4"]}\n\n{output_dict["5"]}'


inputs = [
    gr.inputs.Textbox(label="Prompt", lines=5, default="COVID-19 is"),
    gr.Dropdown(model_list, value="microsoft/biogpt", label="Model ID"),
    gr.inputs.Slider(5, 100, 25, default=25, label="Max Length"),
    gr.inputs.Slider(1, 10, 5, default=5, label="Num Return Sequences")
    ]

outputs = gr.outputs.Textbox(label="Output")

examples = [
    ["COVID-19 is", "microsoft/biogpt"]
]

title = " BioGPT: Generative Pre-trained Transformer for Biomedical Text Generation and Mining"
demo_app = gr.Interface(
    biogpt, 
    inputs, 
    outputs, 
    title=title, 
    examples=examples, 
    cache_examples=True,
)
demo_app.launch(debug=True, enable_queue=True)