import gradio as gr
from transformers import pipeline


article = '''<img src="https://corporateweb-v3-corporatewebv3damstrawebassetbuck-1lruglqypgb84.s3-ap-southeast-2.amazonaws.com/public/cta-2.jpg"/> '''

examples = [
  [
'''
A truck narrowly missed a person on a bicycle when they were reversing out of the depot on Friday.  \
    It was early morning before the sun was up and the cyclist did not have a light.  Fortunately the \
        driver spotted the rider and braked heavily to avoid a collision.
'''],
  [
'''
When making a coffee I noticed the cord to the coffee machine was frayed and tagged it out of service.  Now I need to find a barista!'''],
  [
'''
A worker was using a grinder in a confined space when he became dizzy from the fumes in the area and had to be helped out.  \
The gas monitor he was using was found to be faulty and when the area was assessed with another monitor there was an \
unacceptably high level of CO2 in the area''']]


title = "Incident Prioritisation Tool"
description = "Triage new incidents based on a distilbert-uncased NLP model that has been fine tuned on descriptions of incidents \
    that have been risk rated in the past"

pipe = pipeline("text-classification", model="mrosinski/autotrain-distilbert-risk-ranker-1593356256")

def predict(text):
    # if len(text[0]) > 60:
    preds = pipe(text)[0]
    return preds["label"].title(), f'Confidence Score: {round(preds["score"]*100, 1)}%'
    # else:
    #     return 'Invalid entry', 'Try adding more information to describe the incident'

  
gradio_ui = gr.Interface(
    fn=predict,         
    title=title,
    description=description,
    inputs=[
        gr.inputs.Textbox(lines=5, label="Paste some text here"),
    ],
    outputs=[
        gr.outputs.Textbox(label="Label"),
        gr.outputs.Textbox(label="Score"),
    ],
    examples=examples, 
    article=article 
    
)

gradio_ui.launch(debug=True)