Spaces:
Runtime error
Runtime error
import streamlit as st | |
import torch | |
from transformers import RobertaTokenizer | |
from model import RobertaClass | |
from text_preprocessing import preprocess_text | |
# Load the fine-tuned BERT model | |
model = RobertaClass() | |
model.load_state_dict(torch.load('model.bin', | |
map_location=torch.device('cpu'))) | |
# Load the tokenizer | |
tokenizer = RobertaTokenizer.from_pretrained( | |
'roberta-base', truncation=True, do_lower_case=True) | |
# Define the user interface | |
st.title('ChatGPT detector') | |
text_input = st.text_input('Enter text to classify:', '') | |
submit_button = st.button('Classify') | |
# Define prediction function | |
def predict(text): | |
'''Predicts the label and confidence level of the input text.''' | |
# Preprocess the input text | |
text_preprocessed = preprocess_text(text) | |
# Tokenize the preprocessed text | |
inputs = tokenizer(text_preprocessed, return_tensors='pt', | |
padding=True, truncation=True) | |
inputs.pop('token_type_ids', None) # Remove token_type_ids | |
# Perform inference | |
with torch.no_grad(): | |
outputs = model(input_ids=inputs['input_ids'], | |
attention_mask=inputs['attention_mask']) | |
# Convert output to probabilities and predicted label | |
# probability of positive class | |
predicted_prob = torch.sigmoid(outputs).item() | |
predicted_label = 1 if predicted_prob >= 0.5 else 0 | |
if predicted_label == 0: | |
predicted_prob = 1 - predicted_prob | |
return predicted_label, predicted_prob | |
# Handle user interaction | |
if submit_button: | |
predicted_label, predicted_prob = predict(text_input) | |
# Assuming binary classification | |
labels = ['written by a human', 'generated by ChatGPT'] | |
predicted_category = labels[predicted_label] | |
predicted_prob_percentage = round(predicted_prob * 100, 2) | |
st.write( | |
f"This text was {predicted_category} ({predicted_prob_percentage} % confident)") | |