Spaces:
Sleeping
Sleeping
import streamlit as st | |
from joblib import load | |
from transformers import BertTokenizer, BertForSequenceClassification | |
import torch | |
from tensorflow.keras.models import load_model | |
import tensorflow as tf | |
from tensorflow.keras.preprocessing.text import Tokenizer | |
from tensorflow.keras.preprocessing.sequence import pad_sequences | |
import time | |
from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
model_checkpoint = 'cointegrated/rubert-tiny-toxicity' | |
toxicity_tokenizer = AutoTokenizer.from_pretrained(model_checkpoint) | |
toxicity_model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint) | |
clf = load('my_model_filename.pkl') | |
vectorizer = load('tfidf_vectorizer.pkl') | |
scaler = load('scaler.joblib') | |
tukinazor = load('tokenizer.pkl') | |
rnn_model = load_model('path_to_my_model.h5') | |
bert_model = BertForSequenceClassification.from_pretrained('my_bert_model') | |
tokenizer = BertTokenizer.from_pretrained('bert-base-multilingual-cased') | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
bert_model = bert_model.to(device) | |
labels = ["не токсичный", "оскорбляющий", "непристойный", "угрожающий", "опасный"] | |
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]) | |
else: | |
# Добавленный блок кода | |
result = {} | |
for label, prob in zip(labels, proba): | |
result[label] = prob | |
return result | |
def predict_text(text): | |
sequences = tukinazor.texts_to_sequences([text]) | |
padded_sequences = tf.keras.preprocessing.sequence.pad_sequences(sequences, maxlen=200, padding='post', truncating='post') | |
predictions = rnn_model.predict(padded_sequences) | |
predicted_class = tf.argmax(predictions, axis=-1).numpy()[0] | |
return predicted_class | |
def main(): | |
page_selection = st.sidebar.selectbox("Выберите страницу:", ["Классификация отзывов", "Анализ токсичности"]) | |
if page_selection == "Классификация отзывов": | |
page_reviews_classification() | |
elif page_selection == "Анализ токсичности": | |
page_toxicity_analysis() | |
if __name__ == "__main__": | |
main() | |
def page_reviews_classification(): | |
st.title("Модель классификации отзывов") | |
# Ввод текста | |
user_input = st.text_area("Введите текст отзыва:") | |
if st.button("Классифицировать"): | |
start_time = time.time() | |
user_input_vec = vectorizer.transform([user_input]) | |
sentence_vector_scaled = scaler.transform(user_input_vec) | |
prediction = clf.predict( | |
sentence_vector_scaled) | |
elapsed_time = time.time() - start_time | |
st.write(f"Прогнозируемый класс: {prediction[0]}") | |
st.write(f"Время вычисления: {elapsed_time:.2f} сек.") | |
user_input_rnn = st.text_area("Введите текст отзыва для RNN модели:") | |
if st.button("Классифицировать с RNN"): | |
start_time = time.time() | |
prediction_rnn = predict_text(user_input_rnn) | |
elapsed_time = time.time() - start_time | |
st.write(f"Прогнозируемый класс с RNN: {prediction_rnn}") | |
st.write(f"Время вычисления: {elapsed_time:.2f} сек.") | |
user_input_bert = st.text_area("Введите текст отзыва для BERT:") | |
if st.button("Классифицировать (BERT)"): | |
start_time = time.time() | |
encoding = tokenizer.encode_plus( | |
user_input_bert, | |
add_special_tokens=True, | |
max_length=200, | |
return_token_type_ids=False, | |
padding='max_length', | |
truncation=True, | |
return_attention_mask=True, | |
return_tensors='pt' | |
) | |
input_ids = encoding['input_ids'].to(device) | |
attention_mask = encoding['attention_mask'].to(device) | |
with torch.no_grad(): | |
outputs = bert_model(input_ids=input_ids, attention_mask=attention_mask) | |
predictions = torch.argmax(outputs.logits, dim=1) | |
elapsed_time = time.time() - start_time | |
st.write(f"Прогнозируемый класс (BERT): {predictions.item() + 1}") | |
st.write(f"Время вычисления: {elapsed_time:.2f} сек.") | |
def page_toxicity_analysis(): | |
# Код для анализа токсичности текста с использованием модели cointegrated/rubert-tiny-toxicity | |
user_input_toxicity = st.text_area("Введите текст для оценки токсичности:") | |
if st.button("Оценить токсичность"): | |
start_time = time.time() | |
probs = text2toxicity(user_input_toxicity, aggregate=False) | |
elapsed_time = time.time() - start_time | |
for label, prob in probs.items(): | |
st.write(f"Вероятность {label}: {prob:.4f}") |