import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

model_checkpoint = "hamzamalik11/Biobart_radiology_summarization"  
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 = [ "heart mediastinal contours normal left sided subclavian line position tip distal svc lungs remain clear active disease effusions",
  "prevoid bladder volume cc postvoid bladder volume cc bladder grossly normal appearance"
 ]

description = """
THIS MODEL SUMMARIZE FINDINGS OF RADIOLOGY REPORTS INTO IMPRESSIONS 
<b>Enter a findings of radiology report to see the generated impression!</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)