shukdevdatta123/SQL-Injection-Classifier-Deepseek-R1-Fine-tune

1import gradio as gr2from unsloth import FastLanguageModel3from transformers import AutoTokenizer4 5# Function to load model and tokenizer6def load_model(hf_token):7    try:8        # Initialize the model using ZeroGPU (to run on CPU in Hugging Face environment)9        model_name = "shukdevdatta123/sql_injection_classifier_DeepSeek_R1_fine_tuned_model"10        model, tokenizer = FastLanguageModel.from_pretrained(11            model_name=model_name,12            load_in_4bit=True,13            token=hf_token,14            # use_zero=True,  # Ensure ZeroGPU usage15        )16        return model, tokenizer17    except Exception as e:18        return None, str(e)19 20# Function to predict SQL injection21def predict_sql_injection(query, hf_token):22    model, tokenizer = load_model(hf_token)23    24    if model is None:25        return f"Error loading model: {tokenizer}"26 27    # Prepare the model for inference28    inference_model = FastLanguageModel.for_inference(model)29    30    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"31    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")32 33    # Use the inference model for generation34    outputs = inference_model.generate(35        input_ids=inputs.input_ids,36        attention_mask=inputs.attention_mask,37        max_new_tokens=1000,38        use_cache=True,39    )40    prediction = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]41    return prediction.split("### Classification:\n")[-1].strip()42 43# Gradio UI44def classify_sql_injection(query, hf_token):45    if not hf_token:46        return "Please enter your Hugging Face token."47    48    if not query:49        return "Please enter a SQL query first."50    51    result = predict_sql_injection(query, hf_token)52    return f"Prediction: {result}"53 54# Gradio interface55iface = gr.Interface(56    fn=classify_sql_injection,57    inputs=[58        gr.Textbox(label="SQL Query", placeholder="Enter SQL query here..."),59        gr.Textbox(label="Hugging Face Token", type="password")60    ],61    outputs="text",62    live=True,63    title="SQL Injection Classifier",64    description="Enter an SQL query and your Hugging Face token to classify whether the query is a normal SQL query (0) or a SQL injection attack (1)."65)66 67# Launch the Gradio app68iface.launch()69