# -*- coding: utf-8 -*-
"""
Created on Wed Mar  1 13:48:19 2023

@author: Alexandru Vasile
"""

import os
import gradio as gr
import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

print(f"Is CUDA available: {torch.cuda.is_available()}")
print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")

examples = [['question: Should chest wall irradiation be included after mastectomy and negative node breast cancer? context: This study aims to evaluate local failure patterns in node negative breast cancer patients treated with post-mastectomy radiotherapy including internal mammary chain only. Retrospective analysis of 92 internal or central-breast node-negative tumours with mastectomy and external irradiation of the internal mammary chain at the dose of 50 Gy, from 1994 to 1998. Local recurrence rate was 5 % (five cases). Recurrence sites were the operative scare and chest wall. Factors associated with increased risk of local failure were age<or = 40 years and tumour size greater than 20mm, without statistical significance. answer: Post-mastectomy radiotherapy should be discussed for a sub-group of node-negative patients with predictors factors of local failure such as age<or = 40 years and larger tumour size. target: the answer to the question given the context is'],
            ['question: Do some U.S. states have higher/lower injury mortality rates than others? context: This article examines the hypothesis that the six U.S. states with the highest rates of road traffic deaths (group 1 states) also had above-average rates of other forms of injury such as falling, poisoning, drowning, fire, suffocation, homicide, and suicide, and also for the retail trade and construction industries. The converse, second hypothesis, for the six states with the lowest rates of road traffic deaths (group 2 states) is also examined. Data for these 12 states for the period 1983 to 1995 included nine categories of unintentional and four categories of intentional injury. Seventy-four percent of the group 1 states conformed to the first hypothesis, and 85% of the group 2 states conformed to the second hypothesis. answer: Group 1 states are likely to exhibit above-average rates for most other categories of injury death, whereas group 2 states are even more likely to exhibit below-average rates for most other categories of injury death. target: the answer to the question given the context is']] 

tokenizer = AutoTokenizer.from_pretrained("microsoft/BioGPT-Large-PubMedQA")

model = AutoModelForCausalLM.from_pretrained("microsoft/BioGPT-Large-PubMedQA")

pipe_biogpt = pipeline("text-generation", 
                       model=model, 
                       tokenizer = tokenizer, 
                       device="cuda:0")

title = "BioGPT Q&A Demo"
description = """
NOTE: this model requires the inputs in the following format: `question: <text> context: <text> answer: <text> target: the answer to the question given the context is`.
Check out the [BioGPT-Large-PubMedQA model card](https://huggingface.co/microsoft/biogpt-large-pubmedqa) for more info. 
**Disclaimer:** this demo was made for research purposes only and should not be used for medical purposes.
"""

def inference(text):
  output_biogpt = pipe_biogpt(text, max_length=100)[0]["generated_text"]
  return [
      output_biogpt.split(' ')[-1], 
  ]

io = gr.Interface(
  inference,
  gr.Textbox(lines=3),
  outputs=[
    gr.Textbox(lines=3, label="BioGPT-Large"),
  ],
  title=title,
  description=description,
  examples=examples
)
io.launch()