Update main.py

main.py

@@ -5,7 +5,6 @@ import os
 
 app = Flask(__name__)
 
-
 # Load model and tokenizer
 def load_model():
     # Load saved config and weights
@@ -18,7 +17,6 @@ def load_model():
     model.eval()
     return model
 
-
 # Load components
 try:
     tokenizer = RobertaTokenizer.from_pretrained("./tokenizer_readability")
@@ -27,24 +25,17 @@ try:
 except Exception as e:
     print(f"Error loading model: {str(e)}")
 
-
 @app.route("/")
 def home():
     return request.url
 
-
 @app.route("/predict")
 def predict():
     try:
-        # Get code from request
-        # data = request.get_json()
-        # if "code" not in data:
-        #     return jsonify({"error": "Missing 'code' parameter"}), 400
-        data = request.get_json(force=True, silent=True)
-        if not data or "code" not in data:
-            return jsonify({"error": f"Missing 'code' parameter. data: {data}"}), 400
-
-        code = data["code"]
+        # Get code from URL parameter
+        code = request.args.get("code")
+        if not code:
+            return jsonify({"error": "Missing 'code' URL parameter"}), 400
 
         # Tokenize input
         inputs = tokenizer(
@@ -54,7 +45,6 @@ def predict():
             max_length=512,
             return_tensors='pt'
         )
-        print("here")
 
         # Make prediction
         with torch.no_grad():
@@ -71,6 +61,5 @@ def predict():
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
-
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)