test_desklib.py

"""
Test script for Desklib pre-trained model
"""
import sys
import os

# Fix macOS MPS issues
if sys.platform == "darwin":
    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
    os.environ["OMP_NUM_THREADS"] = "1"
    os.environ["PYTORCH_ENABLE_MPS"] = "0"

import torch
if sys.platform == "darwin":
    try:
        torch.backends.mps.enabled = False
        torch.set_default_device("cpu")
    except:
        pass

from ai_text_detector.models import DetectorModel

print("🧪 Testing Desklib Pre-trained Model")
print("=" * 60)

# Load model
print("\n📥 Loading Desklib model...")
model = DetectorModel("desklib/ai-text-detector-v1.01", use_desklib=True)
print("✅ Model loaded!")

# Test texts
test_texts = [
    ("AI detection refers to the process of identifying whether a given piece of content, such as text, images, or audio, has been generated by artificial intelligence.", "AI"),
    ("I went to the store yesterday and bought some milk and bread. It was a nice sunny day.", "Human"),
]

print("\n🔍 Testing predictions...")
print("=" * 60)

for text, expected in test_texts:
    ai_prob, label = model.predict(text)
    result = "🤖 AI-generated" if label == 1 else "🧑 Human-written"
    print(f"\nText: {text[:80]}...")
    print(f"Prediction: {result}")
    print(f"AI Probability: {ai_prob:.2%}")
    print(f"Expected: {expected}")

print("\n✅ Test complete!")