pipeline

app.py

@@ -1,4 +1,29 @@
 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)
+# 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)
+    score = out['score'] * 4 - 2
+    if score >= 0.5:
+        label = 'not OK'
+    else:
+        label = 'OK'
+    st.json({'label' : label, 'score' : score})

requirements.txt

@@ -0,0 +1,3 @@
+torch
+transformers
+safetensors