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}")