Spaces:
Running
Running
| import streamlit as st | |
| import pandas as pd | |
| import pickle | |
| import transformers as tfs | |
| from catboost import CatBoostClassifier | |
| from sklearn.linear_model import LogisticRegression | |
| import numpy as np | |
| import torch as t | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| st.set_page_config(page_title="# Выявление негативных комментариев с BERT") | |
| st.markdown('# Выявление негативных комментариев с BERT') | |
| with st.expander("Описание проекта"): | |
| st.write(""" | |
| Интернет-магазин «Викишоп» запускает новый сервис. | |
| Теперь пользователи могут редактировать и дополнять описания товаров, как в вики-сообществах. | |
| То есть клиенты предлагают свои правки и комментируют изменения других. | |
| Магазину нужен инструмент, который будет искать токсичные комментарии и отправлять их на модерацию. | |
| """) | |
| def detect_language(text): | |
| first_letter = text[0].lower() | |
| if 'а' <= first_letter <= 'я': | |
| return 'ru' | |
| def ru_bert_comments(text): | |
| model_checkpoint = 'cointegrated/rubert-tiny-toxicity' | |
| tokenizer = AutoTokenizer.from_pretrained(model_checkpoint) | |
| model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint) | |
| if t.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 t.no_grad(): | |
| inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True).to(model.device) | |
| proba = t.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 | |
| return round(text2toxicity(text, True)) | |
| def preprocessing(text): | |
| lang = detect_language(text) | |
| if lang == "ru": | |
| return ru_bert_comments(str(text)) | |
| else: | |
| tokenizer = tfs.AutoTokenizer.from_pretrained('unitary/toxic-bert') | |
| model = tfs.AutoModel.from_pretrained('unitary/toxic-bert') | |
| tokenized = tokenizer.encode(text, add_special_tokens=True) | |
| if len(tokenized) > 512: | |
| truncated_tokens = tokenized[:510] | |
| truncated_tokens = [101] + truncated_tokens + [102] | |
| tokenized = truncated_tokens | |
| padded = np.array(tokenized + [0] * (512 - len(tokenized))) | |
| attention_mask = np.where(padded != 0, 1, 0) | |
| input_ids = t.tensor(padded) | |
| attention_mask = t.tensor(attention_mask) | |
| with t.no_grad(): | |
| embeddings = model(input_ids.unsqueeze(0), attention_mask=attention_mask.unsqueeze(0))[0][:, 0, :].cpu().numpy() | |
| return embeddings | |
| def query(features): | |
| model = pickle.load(open('models/toxic_comments_bert.pkl', 'rb')) | |
| predict = model.predict(features) | |
| return predict | |
| comment = st.text_area("Введите ваш комментарий, модель работает на английском и русском языках и нажмите Ctrl+Enter", "") | |
| result = None | |
| if comment: | |
| st.markdown('## Результат:') | |
| embeddings = preprocessing(comment) | |
| if isinstance(embeddings, int): | |
| if embeddings == 0: | |
| result = 'Комментарий не токсичный' | |
| else: | |
| result = 'Комментарий является токсичным' | |
| else: | |
| if query(embeddings) == 0: | |
| result = 'Комментарий не токсичный' | |
| else: | |
| result = 'Комментарий является токсичным' | |
| if result is not None: | |
| st.write(result) | |