import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification

model_name = "NimaKL/FireWatch_tiny_75k"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

def predict(text):
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits
    label_id = logits.argmax(axis=1).item()
    return "Danger of fire hazard!" if label_id == 1 else "It is unlikely that a fire will start in this area."

# Define a custom CSS style
custom_style = """
    body {
        background-color: #F262626;
    }
"""

# Define a function to generate HTML for embedding the Google Sheets document
def get_sheet_html():
    return f'<iframe src="https://docs.google.com/spreadsheets/d/1Cgq9g2H2yXJSbT0ry8Z92b-6oEQwYS8bCt4ioHQ_gy8/edit?usp=drivesdk" width="640" height="480"></iframe>'

io = gr.Interface(
    fn=predict,
    inputs="text",
    outputs="text",
    title="FireWatch",
    description="<h2>Predict whether a data row describes a fire hazard or not. </h2>\
    <br><br><h3>Here is a <a href='https://docs.google.com/spreadsheets/d/1Cgq9g2H2yXJSbT0ry8Z92b-6oEQwYS8bCt4ioHQ_gy8/preview'>Google Sheets document</a> containing sample data (You can use for testing). It is a heavy document so it might take a while to load.<h3>",
    output_description="Prediction",
    examples=[['-26.76123, 147.15512, 393.02, 203.63'], ['-26.7598, 147.14514, 361.54, 79.4'], ['-25.70059, 149.48932, 313.9, 5.15'], ['-24.4318, 151.83102, 307.98, 8.79'], ['-23.21878, 148.91298, 314.08, 7.4'], ['7.87518, 19.9241, 316.32, 39.63'], ['-20.10942, 148.14326, 314.39, 8.8'], ['7.87772, 19.9048, 304.14, 13.43'], ['-20.79866, 124.46834, 366.74, 89.06']],
    theme="Streamlit",
    css=custom_style
)
io.launch()