# -*- coding: utf-8 -*- Nour Eddine Zekaoui et al.

import gradio as gr
from transformers import pipeline


summarizer = pipeline("summarization", model="NouRed/medqsum-bart-large-xsum-meqsum")

def summarize(
    chq,
    temperature:float=0.1,
    max_length:int=32,
    top_k:float=40,
    top_p:float=0.95,
    do_sample: bool=True,
    num_beams:int=2
    ):
    
    summary = summarizer(
        chq,
        temperature=temperature,
        max_length=max_length,
        top_k=top_k,
        top_p=top_p,
        do_sample=do_sample,
        num_beams=num_beams,
        repetition_penalty=2.5,
        length_penalty=1.0,
        early_stopping=True
        )
    
    return  summary[0]['summary_text']


description = """
<div style="justify-content: center; text-align: center;">
    <h2>
        Enhancing Large Language Models' Utility for Medical Question-Answering: A Patient Health Question Summarization Approach
    </h2>
    <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
        <a href="http://dx.doi.org/10.1109/SITA60746.2023.10373720" target="_blank"><img src="https://img.shields.io/badge/IEEE-10373720-red"></a>
        <a href="https://huggingface.co/NouRed/medqsum-bart-large-xsum-meqsum" target="_blank"><img src="https://img.shields.io/badge/🤗_Hugging_Face-MEDQSUM-orange" alt="HF HUB"></a>
        <a href="https://github.com/zekaouinoureddine/MedQSum" target="_blank"><img src="https://img.shields.io/badge/GitHub-Repo-blue"></a>
    </div>
</div>
"""

gr.Interface(
    fn=summarize,
    inputs=[
        gr.components.Textbox(
            lines=2,
            label="Consumer Health Question",
            placeholder="Type your CHQ ...",
        ),
        gr.components.Slider(
            minimum=0, maximum=1, value=0.1, label="Temperature"
        ),
        gr.components.Slider(
            minimum=1, maximum=2000, step=1, value=32, label="Max tokens"
        ),
       gr.components.Slider(
            minimum=0, maximum=100, step=1, value=40, label="Top k"
        ),
        gr.components.Slider(
            minimum=0, maximum=1, value=0.9, label="Top p"
        ),
        gr.components.Checkbox(
            value=True, label="Do Sample", info="Do you want to use sampling during text generation?"
        ),
        gr.components.Slider(
            minimum=1, maximum=4, step=1, value=2, label="Beams"
        ),
    ],
    outputs=[
        gr.components.Textbox(
            lines=5,
            label="Summary",
        )
    ],
    examples=[
        [" SUBJECT: high inner eye pressure above 21 possible glaucoma  MESSAGE: have seen inner eye pressure increase as I have begin taking Rizatriptan. I understand the med narrows blood vessels. Can this med. cause or effect the closed or wide angle issues with the eyelense/glacoma.",
         0.1,
         32,
         40,
         0.95,
         True,
         2],
        ["Hey I was just wanting to know how I can try to receive stem cell treatment for spinal cord injury using the stem cells I have banked from my newborn baby's umbilical cord tissue? If u have any information please tell me, you'll be helping to save my life!",
         0.1,
         32,
         40,
         0.95,
         True,
         2],

    ],
    theme="soft",
    description=description,
).launch()