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

ACCESS_TOKEN = "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"

app = Flask(__name__)

model_path = "maheshj01/sql-injection-classifier"
model = AutoModelForSequenceClassification.from_pretrained(model_path, token=ACCESS_TOKEN)
tokenizer = AutoTokenizer.from_pretrained(model_path, token=ACCESS_TOKEN)

# Function to classify a SQL query
def classify_query(query):
    inputs = tokenizer(query, return_tensors="pt", truncation=True, padding=True)
    outputs = model(**inputs)
    prediction = outputs.logits.argmax(-1).item()
    return "Vulnerable" if prediction == 1 else "Secure"

# Example usage
query = "SELECT Column"
result = classify_query(query)
print(f"The query is classified as: {result}")