Spaces:
Runtime error
Runtime error
| import torch | |
| import gradio as gr | |
| from transformers import BioGptTokenizer, BioGptForCausalLM | |
| model_names = [ | |
| "BioGPT", | |
| "BioGPT-Large", | |
| "BioGPT-QA-PubMedQA-BioGPT", | |
| "BioGPT-QA-PubMEDQA-BioGPT-Large", | |
| "BioGPT-RE-BC5CDR", | |
| "BioGPT-RE-DDI", | |
| "BioGPT-RE-DTI", | |
| "BioGPT-DC-HoC" | |
| ] | |
| def load_model(model_name="BioGPT"): | |
| model_name_map = { | |
| "BioGPT":"microsoft/biogpt", | |
| "BioGPT-QA-PubMedQA-BioGPT":"microsoft/BioGPT-Large-PubMedQA" | |
| } | |
| tokenizer = BioGptTokenizer.from_pretrained(model_name_map[model_name]) | |
| model = BioGptForCausalLM.from_pretrained(model_name_map[model_name]) | |
| return tokenizer, model | |
| def get_beam_output(sentence, selected_model, min_len=100,max_len=512, n_beams=1): | |
| tokenizer, model = load_model(selected_model) | |
| inputs = tokenizer(sentence, return_tensors="pt") | |
| with torch.no_grad(): | |
| beam_output = model.generate(**inputs, | |
| min_length=min_len, | |
| max_length=max_len, | |
| num_beams=n_beams, | |
| early_stopping=True | |
| ) | |
| output=tokenizer.decode(beam_output[0], skip_special_tokens=True) | |
| return output | |
| inputs = [ | |
| gr.inputs.Textbox(label="prompt", lines=5, default="Bicalutamide"), | |
| gr.Dropdown(model_names, value="BioGPT", label="selected_model"), | |
| gr.inputs.Slider(1, 500, 1, default=100, label="min_len"), | |
| gr.inputs.Slider(1, 2048, 1, default=1024, label="max_len"), | |
| gr.inputs.Slider(1, 10, 1, default=5, label="num_beams") | |
| ] | |
| outputs = gr.outputs.Textbox(label="output") | |
| examples = [ | |
| ["Bicalutamide", "BioGPT", 25, 100, 5], | |
| ["Janus kinase 3 (JAK-3)", "BioGPT", 25, 100, 5], | |
| ["Apricitabine", "BioGPT", 25, 100, 5], | |
| ] | |
| iface = gr.Interface( | |
| fn=get_beam_output, | |
| inputs=inputs, | |
| outputs=outputs, | |
| examples=examples, | |
| title="BioGPT: generative pre-trained transformer for biomedical text generation and mining" | |
| ) | |
| iface.launch(debug=True, enable_queue=True) |