Update main.py

main.py

@@ -1,67 +1,57 @@
 from flask import Flask, request, jsonify
 import torch
-from transformers import RobertaTokenizer, RobertaForSequenceClassification, RobertaConfig
+from transformers import RobertaTokenizer
 import os
+from transformers import RobertaForSequenceClassification
+import torch.serialization
 
+# Initialize Flask app
 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 the trained model and tokenizer
+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 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)}")
+# Ensure the model is in evaluation mode
+model.eval()
 
 @app.route("/")
 def home():
-    return "Code Readability Scoring API - Send POST request to /predict with code snippet"
+    return request.url
 
 @app.route("/predict", methods=["POST"])
 def predict():
     try:
-        # Get code from request
+        # Debugging: print input code to check if the request is received correctly
+        print("Received code:", request.get_json()["code"])
         data = request.get_json()
         if "code" not in data:
             return jsonify({"error": "Missing 'code' parameter"}), 400
             
-        code = data["code"]
+        code_input = data["code"]
         
-        # Tokenize input
+        # Tokenize the input code using the CodeBERT tokenizer
         inputs = tokenizer(
-            code,
+            code_input,
+            return_tensors='pt',
             truncation=True,
             padding='max_length',
-            max_length=512,
-            return_tensors='pt'
+            max_length=512
         )
         
-        # Make prediction
+        # Make prediction using the model
         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
         
-        # 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
-        })
-        
+        return jsonify({"predicted_score": prediction})
     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)