TextClassifier / app.py
ytkoa's picture
Update app.py
dfc9b8c verified
import gradio as gr
import yaml
from src.load_model import model, device
config = yaml.safe_load(open('config.yaml', 'r'))
threshold = config['predct']['threshold']
def predict(text: str):
prediction = model(text).item()
label = "Токсичный❌" if prediction >= threshold else "Не токсичный✅"
return label, float(prediction)
examples = [
["Спасибо за подробный разбор, это действительно полезно!"],
["Интересный подход, я бы добавил ещё пару примеров для наглядности."],
["Никогда не задумывался об этом с такой точки зрения. Подумаю над вашей идеей."],
["папа вроде нормальным был а сынок говнюком вырос."],
["говно на палке блять чё красивого в этой картинке"],
["идиоты! что попало придумывают лишь бы лайки ставили"]
]
interface = gr.Interface(
fn=predict,
title="Text Classification",
description=f"using device: {device}",
inputs=gr.Textbox(label="Введите текст для проверки наличия токсичности"),
outputs=[
gr.Textbox(label="Класс", interactive=False),
gr.Slider(minimum=0, maximum=1, label="Оценка модели", interactive=False)
],
live=True,
examples=examples
)
if __name__ == "__main__":
interface.launch()