salmane11/SQLQueryShield

1---2library_name: transformers3tags:4- text-to-SQL5- SQL6- code-generation7- NLQ-to-SQL8- text2SQL9- Security10- Vulnerability detection11datasets:12- salmane11/SQLShield13language:14- en15base_model:16- microsoft/codebert-base17---18 19# SQLQueryShield20 21## Model Description22 23SQLQueryShield is a vulnerable SQL query detection model. It classifies SQL queries as either vulnerable (e.g., prone to SQL injection or unsafe execution) or benign (safe to execute).24 25The checkpoint included in this repository is based on [microsoft/codebert-base](https://huggingface.co/microsoft/codebert-base) and further finetuned on [SQLShield](https://huggingface.co/datasets/salmane11/SQLShield), a dataset dedicated to text-to-SQL vulnerability detection composed of vulnerable and safe NLQs and their related SQL queries.26 27 28## Finetuning Procedure29The model was fine-tuned using the Hugging Face Transformers library. The following steps were used:30 311. Dataset: SSQLShield, only the SQL queries from the (NLQ, SQL) pairs were used for training.32 332. Preprocessing:34 35    - Input Format: Raw SQL query strings.36 37    - Tokenization: Tokenized using microsoft/codebert-base.38 39    - Max Length: 128 tokens.40 41    - Padding and truncation applied.42 43## Intended Use and Limitations44 45SQLQueryShield is intended for use as a post-generation filter or analysis tool in any system that executes or generates SQL queries. Its main role is to detect whether a SQL query is potentially harmful due to vulnerability patterns such as SQL injection, improper string concatenation, or unsafe expressions.46 47Ideal use cases:48 49    - Filtering SQL queries in Text-to-SQL applications50 51    - Post-processing or validating user-generated SQL before execution52 53 54## How to Use55 56Example 1: Malicious57 58```python59from transformers import pipeline60 61sql_query_shield = pipeline("text-classification", model="salmane11/SQLQueryShield")62 63# For the following Table schema64# CREATE TABLE campuses65#   (66#      campus   VARCHAR,67#      location VARCHAR68#   )69 70query = "SELECT campus FROM campuses WHERE location = '' UNION SELECT database() --"71 72prediction = sql_query_shield(query)73print(prediction)74#[{'label': 'MALICIOUS', 'score': 0.9995294809341431}]75```76 77 78Example 2: Safe79 80```python81from transformers import pipeline82 83sql_query_shield = pipeline("text-classification", model="salmane11/SQLQueryShield")84 85# For the following Table schema86# CREATE TABLE tv_channel87#   (88#      package_option VARCHAR,89#      series_name    VARCHAR90#   ) 91 92query = "SELECT package_option FROM tv_channel WHERE series_name = 'Sky Radio'"93 94 95prediction = sql_query_shield(query)96print(prediction)97#[{'label': 'SAFE', 'score': 0.999503493309021}]98```99 100 101## Cite our work102 103Citation