Update main.py

main.py

@@ -2,13 +2,16 @@ from flask import Flask, request, jsonify
 import torch
 from transformers import RobertaTokenizer
 import os
-
+from transformers import RobertaForSequenceClassification
+import torch.serialization
 # Initialize Flask app
 app = Flask(__name__)
 
 # Load the trained model and tokenizer
 tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")
-model = torch.load("model.pth", map_location=torch.device('cpu'))  # Load the trained model
+torch.serialization.add_safe_globals([RobertaForSequenceClassification])
+
+model = torch.load("model.pth", map_location=torch.device('cpu'), weights_only=False)  # Load the trained model
 
 # Ensure the model is in evaluation mode
 model.eval()
@@ -54,5 +57,5 @@ def predict():
 
 
 # Run the Flask app
-if __name__ == "main":
+if __name__ == "__main__":
     app.run(debug=True)