import streamlit as st
import safetensors
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline

# x = st.slider('Select a value')
# st.write(x, 'squared is', x * x)

name = 'KoalaAI/Text-Moderation'
model = AutoModelForSequenceClassification.from_pretrained(name, num_labels=1, ignore_mismatched_sizes=True)
tokenizer = AutoTokenizer.from_pretrained(name)

d = {}
with safetensors.safe_open("model.safetensors", framework="pt", device='cpu') as f:
    for k in f.keys():
        d[k] = f.get_tensor(k)

model.load_state_dict(d)
pipe = pipeline("text-classification", model=model, tokenizer=tokenizer, device='cpu')

text = st.text_area("enter the text")

if text:
    out = pipe(text)[0]
    score = out['score'] * 4 - 2
    if score >= 0.5:
        label = 'not OK'
    else:
        label = 'OK'
    st.json({'label' : label, 'score' : score})