from flask import Flask, request, jsonify, send_from_directory
import requests
import io
import random
import os
from PIL import Image
from datetime import datetime
import string
import re

app = Flask(__name__)

API_URL = "https://api-inference.huggingface.co/models/mohsenfayyaz/toxicity-classifier"
API_TOKEN = os.getenv("HF_READ_TOKEN")  # it is free
headers = {"Authorization": f"Bearer {API_TOKEN}"}

def query(prompt):
    payload = {
        "inputs": prompt
    }

    response = requests.post(API_URL, headers=headers, json=payload)
    response.raise_for_status()
    return response.json()

@app.route("/generate", methods=["POST"])
def generate():
    try:
        data = request.get_json()
        prompt = data["prompt"]

        resp = query(prompt)
        response = {
            "success": True,
            "result": resp
        }

    except Exception as e:
        response = {
            "success": False,
            "error": str(e)
        }

    return jsonify(response)

if __name__ == "__main__":
    app.run(debug=True, host="0.0.0.0", port=7860)