Update main.py

main.py

@@ -8,7 +8,7 @@ app = Flask(__name__)
 # Load model and tokenizer
 def load_model():
     # Load saved config and weights
-    checkpoint = torch.load("codebert_readability_scorer.pth", map_location=torch.device('cpu'))
+    checkpoint = torch.load("codebert_vulnerability_scorer.pth", map_location=torch.device('cpu'))
     config = RobertaConfig.from_dict(checkpoint['config'])
 
     # Initialize model with loaded config
@@ -19,7 +19,7 @@ def load_model():
 
 # Load components
 try:
-    tokenizer = RobertaTokenizer.from_pretrained("./tokenizer_readability")
+    tokenizer = RobertaTokenizer.from_pretrained("./tokenizer_vulnerability")
     model = load_model()
     print("Model and tokenizer loaded successfully!")
 except Exception as e:
@@ -54,7 +54,7 @@ def predict():
         score = torch.sigmoid(outputs.logits).item()
 
         return jsonify({
-            "readability_score": round(score, 4),
+            "vulnerability_score": round(score, 4),
             "processed_code": code[:500] + "..." if len(code) > 500 else code
         })