Update main.py

main.py

@@ -17,7 +17,7 @@ app = Flask(__name__)
 tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")
 torch.serialization.add_safe_globals([RobertaForSequenceClassification])
 
-model = torch.load("model.pth", map_location=torch.device('cpu'), weights_only=False)  # Load the trained model
+model = torch.load("codebert_readability_scorer.pth", map_location=torch.device('cpu'), weights_only=False)  # Load the trained model
 
 # Ensure the model is in evaluation mode
 model.eval()
@@ -34,7 +34,7 @@ def predict():
     print("Received code:", request.get_json()["code"])
     code = request.get_json()["code"]
     # Load saved weights and config
-    checkpoint = torch.load("codebert_vulnerability_scorer.pth")
+    checkpoint = torch.load("codebert_readability_scorer.pth")
     config = RobertaConfig.from_dict(checkpoint['config'])
 
     # Rebuild the model with correct architecture