Modify request type + take params from JSON body.

main.py

@@ -8,6 +8,7 @@ app = Flask(__name__)
 # Load model and tokenizer
 def load_model():
     # Load saved config and weights
+    # Updated to 'codebert_vulnerability_scorer.pth' as per user's request
     checkpoint = torch.load("codebert_vulnerability_scorer.pth", map_location=torch.device('cpu'))
     config = RobertaConfig.from_dict(checkpoint['config'])
 
@@ -19,6 +20,7 @@ def load_model():
 
 # Load components
 try:
+    # Updated to './tokenizer_vulnerability' as per user's request
     tokenizer = RobertaTokenizer.from_pretrained("./tokenizer_vulnerability")
     model = load_model()
     print("Model and tokenizer loaded successfully!")
@@ -29,13 +31,14 @@ except Exception as e:
 def home():
     return request.url
 
-@app.route("/predict")
+@app.route("/predict", methods=["POST"])  
 def predict():
     try:
-        # Get code from URL parameter
-        code = request.args.get("code")
-        if not code:
-            return jsonify({"error": "Missing 'code' URL parameter"}), 400
+        data = request.get_json()
+        if not data or "code" not in data:
+            return jsonify({"error": "Missing 'code' in request body"}), 400
+
+        code = data["code"]
 
         # Tokenize input
         inputs = tokenizer(
@@ -54,8 +57,7 @@ def predict():
         score = torch.sigmoid(outputs.logits).item()
 
         return jsonify({
-            "score": round(score, 4),
-            "given_code": code[:500] + "..." if len(code) > 500 else code
+            "score": round(score, 4)
         })
 
     except Exception as e: