from flask import Flask, request, jsonify, render_template_string
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
app = Flask(__name__)
print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained("Redfire-1234/bert-ai-human-model")
model = AutoModelForSequenceClassification.from_pretrained("Redfire-1234/bert-ai-human-model")
model.eval()
print("Model loaded!")
def predict_text(text):
"""Predict whether text is AI or Human generated"""
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
probs = torch.softmax(logits, dim=1).numpy()[0]
predicted_class = int(torch.argmax(logits, dim=1))
label_map = {0: "Human", 1: "AI"}
return {
"label": label_map[predicted_class],
"confidence": float(probs[predicted_class]),
"probabilities": {"human": float(probs[0]), "ai": float(probs[1])}
}
HTML_TEMPLATE = """
AI vs Human Text Classifier
🤖 AI vs Human Text Classifier
Enter text below to check if it was written by a human or AI
Analyzing...
Confidence:
Probabilities:
Human:
AI:
"""
@app.route('/')
def home():
return render_template_string(HTML_TEMPLATE)
@app.route('/predict', methods=['POST'])
def predict():
try:
data = request.get_json()
if not data or 'text' not in data:
return jsonify({'error': 'No text provided'}), 400
text = data['text']
if not text.strip():
return jsonify({'error': 'Text cannot be empty'}), 400
result = predict_text(text)
return jsonify(result)
except Exception as e:
print(f"Error: {e}")
return jsonify({'error': str(e)}), 500
@app.route('/health')
def health():
return jsonify({'status': 'healthy'})
if __name__ == '__main__':
app.run(host='0.0.0.0', port=7860)