ElPremOoO/Code_Mate

1from flask import Flask, request, jsonify2import torch3from transformers import RobertaTokenizer4import os5from transformers import RobertaForSequenceClassification6import torch.serialization7 8# Initialize Flask app9app = Flask(__name__)10 11# Load the trained model and tokenizer12tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")13torch.serialization.add_safe_globals([RobertaForSequenceClassification])14model = torch.load("model.pth", map_location=torch.device('cpu'), weights_only=False)15 16# Ensure the model is in evaluation mode17model.eval()18 19@app.route("/")20def home():21    return request.url22 23@app.route("/predict")24def predict():25    try:26        # Debugging: print input code to check if the request is received correctly27        print("Received code:", request.get_json()["code"])28        data = request.get_json()29        if "code" not in data:30            return jsonify({"error": "Missing 'code' parameter"}), 40031            32        code_input = data["code"]33        34        # Tokenize the input code using the CodeBERT tokenizer35        inputs = tokenizer(36            code_input,37            return_tensors='pt',38            truncation=True,39            padding='max_length',40            max_length=51241        )42        43        # Make prediction using the model44        with torch.no_grad():45            outputs = model(**inputs)46            prediction = outputs.logits.squeeze().item()47            48        # Extract the predicted score (single float)49        print(f"Predicted score: {prediction}")  # Debugging: Print prediction50        51        return jsonify({"predicted_score": prediction})52    except Exception as e:53        return jsonify({"error": str(e)}), 50054 55# Run the Flask app56if __name__ == "__main__":57    app.run(host="0.0.0.0", port=7860)