ElPremOoO commited on
Commit
c0734d1
·
verified ·
1 Parent(s): 175968e

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +23 -24
main.py CHANGED
@@ -25,37 +25,36 @@ def home():
25
  # @app.route("/predict", methods=["POST"])
26
  @app.route("/predict")
27
  def predict():
28
- return "predicted brooo"
29
- # try:
30
- # # Debugging: print input code to check if the request is received correctly
31
- # print("Received code:", request.get_json()["code"])
32
 
33
- # data = request.get_json()
34
- # if "code" not in data:
35
- # return jsonify({"error": "Missing 'code' parameter"}), 400
36
 
37
- # code_input = data["code"]
38
 
39
- # # Tokenize the input code using the CodeBERT tokenizer
40
- # inputs = tokenizer(
41
- # code_input,
42
- # return_tensors='pt',
43
- # truncation=True,
44
- # padding='max_length',
45
- # max_length=512
46
- # )
47
 
48
- # # Make prediction using the model
49
- # with torch.no_grad():
50
- # outputs = model(**inputs)
51
- # prediction = outputs.logits.squeeze().item() # Extract the predicted score (single float)
52
 
53
- # print(f"Predicted score: {prediction}") # Debugging: Print prediction
54
 
55
- # return jsonify({"predicted_score": prediction})
56
 
57
- # except Exception as e:
58
- # return jsonify({"error": str(e)}), 500
59
 
60
 
61
  # Run the Flask app
 
25
  # @app.route("/predict", methods=["POST"])
26
  @app.route("/predict")
27
  def predict():
28
+ try:
29
+ # Debugging: print input code to check if the request is received correctly
30
+ print("Received code:", request.get_json()["code"])
 
31
 
32
+ data = request.get_json()
33
+ if "code" not in data:
34
+ return jsonify({"error": "Missing 'code' parameter"}), 400
35
 
36
+ code_input = data["code"]
37
 
38
+ # Tokenize the input code using the CodeBERT tokenizer
39
+ inputs = tokenizer(
40
+ code_input,
41
+ return_tensors='pt',
42
+ truncation=True,
43
+ padding='max_length',
44
+ max_length=512
45
+ )
46
 
47
+ # Make prediction using the model
48
+ with torch.no_grad():
49
+ outputs = model(**inputs)
50
+ prediction = outputs.logits.squeeze().item() # Extract the predicted score (single float)
51
 
52
+ print(f"Predicted score: {prediction}") # Debugging: Print prediction
53
 
54
+ return jsonify({"predicted_score": prediction})
55
 
56
+ except Exception as e:
57
+ return jsonify({"error": str(e)}), 500
58
 
59
 
60
  # Run the Flask app