Update main.py

main.py

@@ -4,12 +4,6 @@ from transformers import RobertaTokenizer
 import os
 from transformers import RobertaForSequenceClassification
 import torch.serialization
-import torch
-from transformers import RobertaTokenizer, RobertaForSequenceClassification, Trainer, TrainingArguments
-from torch.utils.data import Dataset
-import pandas as pd
-from sklearn.model_selection import train_test_split
-import numpy as np
 # Initialize Flask app
 app = Flask(__name__)
 
@@ -31,39 +25,38 @@ def home():
 # @app.route("/predict", methods=["POST"])
 @app.route("/predict")
 def predict():
-    print("Received code:", request.get_json()["code"])
-    code = request.get_json()["code"]
-    # Load saved weights and config
-    checkpoint = torch.load("codebert_readability_scorer.pth")
-    config = RobertaConfig.from_dict(checkpoint['config'])
-
-    # Rebuild the model with correct architecture
-    model = RobertaForSequenceClassification(config)
-    model.load_state_dict(checkpoint['model_state_dict'])
-    model.eval()
-
-    # Load tokenizer
-    tokenizer = RobertaTokenizer.from_pretrained('./tokenizer')
-
-    # Prepare input
-    inputs = tokenizer(
-        code,
-        truncation=True,
-        padding='max_length',
-        max_length=512,
-        return_tensors='pt'
-    )
-
-    # Make prediction
-    with torch.no_grad():
-        outputs = model(**inputs)
-
-    score = torch.sigmoid(outputs.logits).item()
-    return score
-    
-
-
-    
+    try:
+        # 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_input = data["code"]
+
+        # Tokenize the input code using the CodeBERT tokenizer
+        inputs = tokenizer(
+            code_input,
+            return_tensors='pt',
+            truncation=True,
+            padding='max_length',
+            max_length=512
+        )
+
+        # 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
+
+        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)