from flask import Flask, render_template, request, jsonify
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

app = Flask(__name__)

MODEL_PATH = "AI-MODEL-FINGERPRINTING/notebooks/saved_distilbert_model"

print("Loading model...")

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)

model.eval()

MODEL_NAMES = [
    "claude",
    "gemini",
    "groq",
    "mistral",
    "openai"
]


@app.route("/")
def home():
    return render_template("index.html")


@app.route("/predict", methods=["POST"])
def predict():

    data = request.get_json()

    text = data.get("text", "").strip()

    if not text:
        return jsonify({"error": "No text provided"})

    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        padding=True,
        max_length=512
    )

    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=1)[0]

    pred_idx = torch.argmax(probs).item()

    prediction = MODEL_NAMES[pred_idx]

    scores = {
        MODEL_NAMES[i]: round(float(probs[i]), 4)
        for i in range(len(MODEL_NAMES))
    }

    return jsonify({
        "prediction": prediction,
        "scores": scores
    })


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