Update app.py
Browse files
app.py
CHANGED
|
@@ -28,7 +28,7 @@ def predict_sql_injection(query, hf_token):
|
|
| 28 |
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("
|
| 32 |
|
| 33 |
# Use the inference model for generation
|
| 34 |
outputs = inference_model.generate(
|
|
|
|
| 28 |
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("cpu")
|
| 32 |
|
| 33 |
# Use the inference model for generation
|
| 34 |
outputs = inference_model.generate(
|