Create main.py

main.py

@@ -0,0 +1,58 @@
+from flask import Flask, request, jsonify
+import torch
+from transformers import RobertaTokenizer
+import os
+
+# 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
+
+Ensure the model is in evaluation mode
+model.eval()
+
+
+@app.route("/")
+def home():
+    return "Model is running!"
+
+
+@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"])
+
+        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(debug=True)