Update main.py

main.py

@@ -90,42 +90,15 @@ except Exception as e:
 @app.route("/", methods=['GET'])
 def home():
     return jsonify({
-        "message": "CodeBERT Vulnerability Scorer API",
+        "message": "CodeBERT Vulnerability Evalutor API",
         "status": "Model loaded" if model is not None else "Model not loaded",
         "device": str(device) if device else "unknown",
         "endpoints": {
-            "/predict": "POST with JSON body containing 'code' field",
-            "/predict_batch": "POST with JSON body containing 'codes' array",
-            "/predict_get": "GET with 'code' URL parameter"
+            "/predict": "POST with JSON body containing 'codes' array"
         }
     })
 
 @app.route("/predict", methods=['POST'])
-def predict_post():
-    try:
-        if model is None or tokenizer is None:
-            return jsonify({"error": "Model not loaded properly"}), 500
-            
-        data = request.get_json()
-        if not data or 'code' not in data:
-            return jsonify({"error": "Missing 'code' field in JSON body"}), 400
-            
-        code = data['code']
-        if not code or not isinstance(code, str):
-            return jsonify({"error": "'code' field must be a non-empty string"}), 400
-            
-        score = predict_vulnerability(code)
-        
-        return jsonify({
-            "score": score,
-            "vulnerability_level": get_vulnerability_level(score),
-            "code_preview": code[:200] + "..." if len(code) > 200 else code
-        })
-        
-    except Exception as e:
-        return jsonify({"error": f"Prediction error: {str(e)}"}), 500
-
-@app.route("/predict_batch", methods=['POST'])
 def predict_batch():
     try:
         if model is None or tokenizer is None:
@@ -148,10 +121,7 @@ def predict_batch():
             
             for j, score in enumerate(scores):
                 results.append({
-                    "index": i + j,
-                    "score": score,
-                    "vulnerability_level": get_vulnerability_level(score),
-                    "code_preview": batch[j][:100] + "..." if len(batch[j]) > 100 else batch[j]
+                    "score": score
                 })
         
         return jsonify({"results": results})
@@ -159,26 +129,7 @@ def predict_batch():
     except Exception as e:
         return jsonify({"error": f"Batch prediction error: {str(e)}"}), 500
 
-@app.route("/predict_get", methods=['GET'])
-def predict_get():
-    try:
-        if model is None or tokenizer is None:
-            return jsonify({"error": "Model not loaded properly"}), 500
-            
-        code = request.args.get("code")
-        if not code:
-            return jsonify({"error": "Missing 'code' URL parameter"}), 400
-            
-        score = predict_vulnerability(code)
-        
-        return jsonify({
-            "score": score,
-            "vulnerability_level": get_vulnerability_level(score),
-            "code_preview": code[:200] + "..." if len(code) > 200 else code
-        })
-        
-    except Exception as e:
-        return jsonify({"error": f"Prediction error: {str(e)}"}), 500
+
 
 def predict_vulnerability(code):
     dynamic_length = min(max(len(code.split()) * 2, 128), 512)
@@ -229,13 +180,6 @@ def predict_vulnerability_batch(codes):
     
     return [round(float(score), 4) for score in scores.flatten()]
 
-def get_vulnerability_level(score):
-    if score < 0.3:
-        return "Low"
-    elif score < 0.7:
-        return "Medium"
-    else:
-        return "High"
 
 @app.route("/health", methods=['GET'])
 def health_check():