|
import gradio as gr |
|
from transformers import TFDistilBertForSequenceClassification, DistilBertTokenizerFast |
|
import tensorflow as tf |
|
|
|
|
|
model_name = "Buebito/HamOrSpam_Model" |
|
tokenizer = DistilBertTokenizerFast.from_pretrained(model_name) |
|
model = TFDistilBertForSequenceClassification.from_pretrained(model_name) |
|
|
|
def classify_text(text): |
|
|
|
inputs = tokenizer(text, return_tensors="tf", truncation=True, padding=True, max_length=512) |
|
|
|
|
|
predictions = model(inputs.data)[0] |
|
|
|
|
|
probabilities = tf.nn.softmax(predictions, axis=-1) |
|
|
|
|
|
prediction_index = tf.argmax(probabilities, axis=-1).numpy()[0] |
|
|
|
|
|
label = "ham" if prediction_index == 0 else "spam" |
|
|
|
|
|
return {"label": label, "probabilities": probabilities.numpy().tolist()} |
|
|
|
|
|
iface = gr.Interface( |
|
fn=classify_text, |
|
inputs=gr.Textbox(lines=2, placeholder="Enter Text Here..."), |
|
outputs=[ |
|
gr.Label(label="Classification"), |
|
gr.JSON(label="Probabilities") |
|
] |
|
) |
|
|
|
|
|
iface.launch() |
|
|