File size: 1,219 Bytes
cb16128 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
import numpy as np
import pandas as pd
# Load the model and tokenizer
model_name = "PL-RnD/privacy-moderation-large"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
# Example text
texts = [
"Here is my credit card number: 1234-5678-9012-3456",
"This is a regular message without sensitive information.",
"For homeowners insurance, select deductibles from $500 to $2,500. Higher deductibles lower premiums.",
"Solidarity: My enrollment includes my kid's braces at $4,000 total—family strained. Push for orthodontic expansions. Email blast to reps starting now.",
]
# Tokenize the input
inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True)
# Get model predictions
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
predictions = torch.argmax(logits, dim=-1)
# Convert predictions to labels
labels = ["non-violation", "violation"]
predicted_labels = [labels[pred] for pred in predictions.numpy()]
# Display results
df = pd.DataFrame({"text": texts, "label": predicted_labels})
print(df)
|