22 lines
812 B
Python
22 lines
812 B
Python
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}") |