Update app.py

app.py

@@ -1,7 +1,68 @@
 import gradio as gr
+from unsloth import FastLanguageModel
+from transformers import AutoTokenizer
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Function to load model and tokenizer
+def load_model(hf_token):
+    try:
+        # Initialize the model using ZeroGPU (to run on CPU in Hugging Face environment)
+        model_name = "shukdevdatta123/sql_injection_classifier_DeepSeek_R1_fine_tuned_model"
+        model, tokenizer = FastLanguageModel.from_pretrained(
+            model_name=model_name,
+            load_in_4bit=True,
+            token=hf_token,
+            use_zero=True,  # Ensure ZeroGPU usage
+        )
+        return model, tokenizer
+    except Exception as e:
+        return None, str(e)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# Function to predict SQL injection
+def predict_sql_injection(query, hf_token):
+    model, tokenizer = load_model(hf_token)
+    
+    if model is None:
+        return f"Error loading model: {tokenizer}"
+
+    # 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()
+
+# Gradio UI
+def classify_sql_injection(query, hf_token):
+    if not hf_token:
+        return "Please enter your Hugging Face token."
+    
+    if not query:
+        return "Please enter a SQL query first."
+    
+    result = predict_sql_injection(query, hf_token)
+    return f"Prediction: {result}"
+
+# Gradio interface
+iface = gr.Interface(
+    fn=classify_sql_injection,
+    inputs=[
+        gr.Textbox(label="SQL Query", placeholder="Enter SQL query here..."),
+        gr.Textbox(label="Hugging Face Token", type="password")
+    ],
+    outputs="text",
+    live=True,
+    title="SQL Injection Classifier",
+    description="Enter an SQL query and your Hugging Face token to classify whether the query is a normal SQL query or a SQL injection attack."
+)
+
+# Launch the Gradio app
+iface.launch()