File size: 1,279 Bytes
37ea438 |
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 32 33 34 35 36 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
def classify_text(text, model_path=None):
# Load model and tokenizer
model_path = model_path or "."
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)
# Prepare the text for the model
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512, padding=True)
# Run inference
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
# Get the predicted class and probabilities
probabilities = torch.nn.functional.softmax(logits, dim=1)
predicted_class_idx = torch.argmax(probabilities, dim=1).item()
confidence = probabilities[0][predicted_class_idx].item()
# Map class index to label
labels = ["Human-written", "AI-generated"]
predicted_label = labels[predicted_class_idx]
return predicted_label, confidence
if __name__ == "__main__":
# Example usage
text = "Enter your text here to test if it's AI-generated or human-written."
result, confidence = classify_text(text)
print(f"This text appears to be: {result}")
print(f"Confidence: {confidence:.4f}")
|