SQLi / Prediction.py
Jamesdagless's picture
Create Prediction.py
e751588
raw
history blame contribute delete
661 Bytes
import joblib
# Load the saved model
model_filename = 'sql_classifier_model.joblib'
loaded_classifier = joblib.load(model_filename)
# Prepare input data for prediction
input_data = ["SELECT * FROM employees", "DELETE FROM users WHERE user_id = 123"]
# Text vectorization (using the same TF-IDF vectorizer)
input_data_tfidf = tfidf_vectorizer.transform(input_data)
# Make predictions using the loaded model
predictions = loaded_classifier.predict(input_data_tfidf)
# Display the predictions
for query, prediction in zip(input_data, predictions):
label = "malicious" if prediction == 1 else "safe"
print(f"SQL Query: {query}\nPrediction: {label}\n")