Update README.md

README.md

@@ -35,25 +35,56 @@ library_name: peft
 
 ## Uses
 
-<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
-
 ### Direct Use
 
-<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+To use the SQL Injection Classifier model, you can follow the code snippet below. This example demonstrates how to predict whether a given SQL query is normal or an injection attack.
 
-[More Information Needed]
+```python
+from unsloth import FastLanguageModel
+from transformers import AutoTokenizer
 
-### Downstream Use [optional]
+# Load the model and tokenizer
+model_name = "unsloth/DeepSeek-R1-Distill-Llama-8B"
+hf_token = "your hf tokens"
 
-<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name=model_name,
+    load_in_4bit=True,
+    token=hf_token,
+)
 
-[More Information Needed]
+# Function for testing queries
+def predict_sql_injection(query):
+    # Prepare the model for inference
+    inference_model = FastLanguageModel.for_inference(model)
+
+    prompt = f"### Instruction:\nClassify the following SQL query as normal (0) or an injection attack (1).\n\n### Query:\n{query}\n\n### Classification:\n"
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+
+    # Use the inference model for generation
+    outputs = inference_model.generate(
+        input_ids=inputs.input_ids,
+        attention_mask=inputs.attention_mask,
+        max_new_tokens=1000,
+        use_cache=True,
+    )
+    prediction = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+    return prediction.split("### Classification:\n")[-1].strip()
+
+# Example usage
+test_query = "SELECT * FROM users WHERE id = '1' OR '1'='1' --"
+result = predict_sql_injection(test_query)
+print(f"Query: {test_query}\nPrediction: {result}")
+```
+
+### Downstream Use
+
+This model can be integrated into applications requiring SQL injection detection, such as web application firewalls, database query analyzers, and security auditing tools. It can help identify and prevent potential vulnerabilities in SQL queries.
 
 ### Out-of-Scope Use
 
-<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+This model should not be used for malicious purposes, such as testing vulnerabilities on unauthorized systems, or for making security decisions without human oversight. It is essential to understand that the model's predictions should be interpreted with caution and supplemented with additional security measures.
 
-[More Information Needed]
 
 ## Bias, Risks, and Limitations