Create main.py

main.py

@@ -0,0 +1,65 @@
+from flask import Flask, request, jsonify
+import torch
+from transformers import RobertaTokenizer, RobertaForSequenceClassification, RobertaConfig
+import os
+
+app = Flask(__name__)
+
+# Load model and tokenizer
+def load_model():
+    # Load saved config and weights
+    checkpoint = torch.load("codebert_readability_scorer.pth", map_location=torch.device('cpu'))
+    config = RobertaConfig.from_dict(checkpoint['config'])
+
+    # Initialize model with loaded config
+    model = RobertaForSequenceClassification(config)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+    return model
+
+# Load components
+try:
+    tokenizer = RobertaTokenizer.from_pretrained("./tokenizer_readability")
+    model = load_model()
+    print("Model and tokenizer loaded successfully!")
+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 URL parameter
+        code = request.args.get("code")
+        if not code:
+            return jsonify({"error": "Missing 'code' URL parameter"}), 400
+
+        # Tokenize input
+        inputs = tokenizer(
+            code,
+            truncation=True,
+            padding='max_length',
+            max_length=512,
+            return_tensors='pt'
+        )
+
+        # Make prediction
+        with torch.no_grad():
+            outputs = model(**inputs)
+
+        # Apply sigmoid and format score
+        score = torch.sigmoid(outputs.logits).item()
+
+        return jsonify({
+            "readability_score": round(score, 4),
+            "processed_code": code[:500] + "..." if len(code) > 500 else code
+        })
+
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=7860)