toxicity-logreg-vs-bert / compare_models.py
santsant93's picture
Upload 5 files
0fd919c verified
Raw
History Blame Contribute Delete
870 Bytes
import pandas as pd
import joblib
from transformers import pipeline
# Load data
sample = pd.read_csv("sample_20.csv")
texts = sample["comment_text"]
true = sample["toxic"].values
# Load Logistic Regression model
vectorizer = joblib.load("tfidf_vectorizer.joblib")
clf = joblib.load("toxic_classifier.joblib")
lr_preds = clf.predict(vectorizer.transform(texts))
# Load BERT model
bert_clf = pipeline(
"text-classification",
model="unitary/toxic-bert",
tokenizer="unitary/toxic-bert",
truncation=True
)
bert_raw = [bert_clf(t)[0]["label"] for t in texts]
bert_preds = [1 if lbl == "toxic" else 0 for lbl in bert_raw]
# Simple accuracy for 20 samples
lr_acc = (lr_preds == true).mean()
bert_acc = (pd.Series(bert_preds).values == true).mean()
print("Logistic Regression accuracy on 20 samples:", lr_acc)
print("BERT accuracy on 20 samples:", bert_acc)