Update main.py

main.py

@@ -1,62 +1,67 @@
 from flask import Flask, request, jsonify
 import torch
-from transformers import RobertaTokenizer
+from transformers import RobertaTokenizer, RobertaForSequenceClassification, RobertaConfig
 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")
-torch.serialization.add_safe_globals([RobertaForSequenceClassification])
 
-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()
+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'))
+    config = RobertaConfig.from_dict(checkpoint['config'])
+    
+    # Initialize model with loaded config
+    model = RobertaForSequenceClassification(config)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+    return model
+
+# Load components
+try:
+    tokenizer = RobertaTokenizer.from_pretrained("./tokenizer")
+    model = load_model()
+    print("Model and tokenizer loaded successfully!")
+except Exception as e:
+    print(f"Error loading model: {str(e)}")
 
 @app.route("/")
 def home():
     return request.url
 
-
-# @app.route("/predict", methods=["POST"])
-@app.route("/predict")
+@app.route("/predict", methods=["POST"])
 def predict():
     try:
-        # Debugging: print input code to check if the request is received correctly
-        print("Received code:", request.get_json()["code"])
-
+        # Get code from request
         data = request.get_json()
         if "code" not in data:
             return jsonify({"error": "Missing 'code' parameter"}), 400
-
-        code_input = data["code"]
-
-        # Tokenize the input code using the CodeBERT tokenizer
+            
+        code = data["code"]
+        
+        # Tokenize input
         inputs = tokenizer(
-            code_input,
-            return_tensors='pt',
+            code,
             truncation=True,
             padding='max_length',
-            max_length=512
+            max_length=512,
+            return_tensors='pt'
         )
-
-        # Make prediction using the model
+        
+        # Make prediction
         with torch.no_grad():
             outputs = model(**inputs)
-            prediction = outputs.logits.squeeze().item()  # Extract the predicted score (single float)
-
-        print(f"Predicted score: {prediction}")  # Debugging: Print prediction
-
-        return jsonify({"predicted_score": prediction})
-
+        
+        # Apply sigmoid and format score
+        score = torch.sigmoid(outputs.logits).item()
+        
+        return jsonify({
+            "readability_score": round(score, 4),
+            "processed_code": code[:500] + "..." if len(code) > 500 else code
+        })
+        
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
-
-# Run the Flask app
 if __name__ == "__main__":
-     app.run(host="0.0.0.0", port=7860)
+    app.run(host="0.0.0.0", port=7860)