Spaces:
Runtime error
Runtime error
| import pandas as pd | |
| import streamlit as st | |
| import torch | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| model_checkpoint = 'cointegrated/rubert-tiny-toxicity' | |
| tokenizer = AutoTokenizer.from_pretrained(model_checkpoint) | |
| model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint) | |
| if torch.cuda.is_available(): | |
| model.cuda() | |
| def text2toxicity(text, aggregate=True): | |
| """ Calculate toxicity of a text (if aggregate=True) or a vector of toxicity aspects (if aggregate=False)""" | |
| with torch.no_grad(): | |
| inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True).to(model.device) | |
| proba = torch.sigmoid(model(**inputs).logits).cpu().numpy() | |
| if isinstance(text, str): | |
| proba = proba[0] | |
| if aggregate: | |
| return 1 - proba.T[0] * (1 - proba.T[-1]) | |
| return proba | |
| text = st.text_area('Введите текст', value='Пороть надо таких придурков!') | |
| proba = text2toxicity(text, aggregate=False) | |
| s = pd.Series( | |
| proba.tolist() + [proba[0] * (1 - proba[-1])], | |
| index=[ | |
| 'Стиль НЕтоксичный', | |
| 'Есть оскорбление', | |
| 'Есть непотребство', | |
| 'Есть угроза', | |
| 'Смысл текста неприемлемый', | |
| 'Текст - ОК' | |
| ], | |
| name='Оценка вероятности' | |
| ) | |
| s | |