Update README.md

README.md

@@ -32,7 +32,7 @@ This model is trained to identify SQL injection attacks, which are a type of cod
 
 ## Uses
 
-### Direct Use
+### Colab Use
 
 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.
 
@@ -95,44 +95,67 @@ This model was trained on a dataset of SQL queries and may exhibit certain limit
 
 Users (both direct and downstream) should be aware of the potential risks of relying on the model in security-sensitive applications. Additional domain-specific testing and validation are recommended before deployment.
 
-## How to Get Started with the Model
+## How to Get Started with the Model (Colab Streamlit)
 
 ```python
+!pip install unsloth
+
+%%writefile app.py
+import streamlit as st
 from unsloth import FastLanguageModel
 from transformers import AutoTokenizer
 
-# Load the model and tokenizer
+# Streamlit UI for input
+st.title("SQL Injection Classifier")
+hf_token = st.text_input("Enter your Hugging Face Token", type="password")
 model_name = "shukdevdatta123/sql_injection_classifier_DeepSeek_R1_fine_tuned_model"
-hf_token = "your hf tokens"
-
-model, tokenizer = FastLanguageModel.from_pretrained(
-    model_name=model_name,
-    load_in_4bit=True,
-    token=hf_token,
-)
-
-# 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}")
+# Load the model and tokenizer when HF token is provided
+if hf_token:
+    try:
+        model, tokenizer = FastLanguageModel.from_pretrained(
+            model_name=model_name,
+            load_in_4bit=True,
+            token=hf_token,
+        )
+
+        # 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()
+
+        # Input query from the user
+        query = st.text_area("Enter an SQL query to test for injection", "")
+        
+        # Add a button to classify the query
+        if st.button("Classify SQL Injection"):
+            if query:
+                result = predict_sql_injection(query)
+                st.write(f"Prediction: {result}")
+            else:
+                st.write("Please enter a SQL query first.")
+
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+else:
+    st.write("Please enter your Hugging Face token to proceed.")
+
+!pip install streamlit
+
+!streamlit run app.py & npx localtunnel --port 8501
 ```
 ## Training Details