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 large screen was being lifted into the plant when the crane operator observed someone walk directly under the suspended load.  \
     The drop zone underneath the load had not been adequately barricaded
     '''
     ],
  ['''
   While building scaffold to access a pipe for repairs a worker was observed not wearing fall protection while the scaffold \
    was still being constructed.  There was a potential for him to fall over 5 metres from the side of the plant
    '''
    ],
  ['''
    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
     '''
     ],
  [
    '''
    Henry Winkler tripped over some rocks that had overflowed from the discharge chute of screen 1011. He had been walking past the screen \
      on the way to the control room.  He suffered a bruised hip and broken wrist and had to be transported to Mackay Base Hospital \
        for treatment.
        '''
        ],
  [
    '''
    Susan Lauper experienced a mild electrical shock while operating an arc welder in the plant today.  It was discovered that the earth \
      was not effectively grounded due to corrosion and buildup in the work area.  She was treated according to site procedures and was \
        given a medical clearance to return to work by her GP later that day.
        '''
        ],
  [
    '''
    While conducting regular inspections of the ground floor, guarding was noticed to be missing on the main drive of pump 1330. \
      This pump had been serviced during the previous maintenance day and it appears that the guard was not replaced prior to startup \
        and was found nearby.  This is a potential breach of isolation procedure and requires further investigation.
        '''
        ]
]


title = "Safety Hazard Classifier"
description = "Using zero shot classification to determine which critical hazard an incident belongs to"

classifier = pipeline("zero-shot-classification", model="Narsil/deberta-large-mnli-zero-cls")

def predict(text):
    preds = classifier(text, candidate_labels=["electrical", "confined space", "unguarded machinery", \
      "spills and tripping hazards", "working from heights", "suspended loads", "machinery related"])
    return dict(zip(preds['labels'], preds['scores']))
  
gradio_ui = gr.Interface(
    fn=predict,         
    title=title,
    description=description,
    inputs=[
        gr.inputs.Textbox(lines=5, label="Paste some text here"),
    ],
    outputs=[
        gr.outputs.Label(num_top_classes=3)
    ],
    examples=examples, 
    article=article 
    
)

gradio_ui.launch(debug=True)