import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

model_checkpoint = "Mbilal755/Radiology_Bart"
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

from transformers import SummarizationPipeline
summarizer = SummarizationPipeline(model=model, tokenizer=tokenizer)   

import gradio as gr

examples = [
  "prevoid bladder volume cc postvoid bladder volume cc bladder grossly normal appearance ",
  "heart mediastinal contours normal left sided subclavian line position tip distal svc lungs remain clear active disease effusions",
  "heart size normal mediastinal hilar contours remain stable small right pneumothorax remains unchanged surgical lung staples overlying left upper lobe seen linear pattern consistent prior upper lobe resection soft tissue osseous structures appear unremarkable nasogastric endotracheal tubes remain satisfactory position atelectatic changes right lower lung field remain unchanged prior study" 
]

description = """
We fine-tuned the BioBart 442M parameter model on a dataset of 52,000 radiology reports for training and 8000 reports for evaluation scraped from MIMIC-III specifically for the task of summarization.
The model is able to generate impressions summarizing key findings from the longer radiology reports. 

<b>Enter a radiology report to see the generated impression summary!</b>
"""

def summarize(radiology_report):
  summary = summarizer(radiology_report)[0]['summary_text']
  return summary

iface = gr.Interface(fn=summarize, 
             inputs=gr.inputs.Textbox(lines=5, label="Radiology Report"),
             outputs=gr.outputs.Textbox(label="Summary"),
             examples=examples,
             title="Radiology Report Summarization",
             description=description,
             theme="huggingface")
             
if __name__ == "__main__":
    iface.launch(share=False)