Update main.py

main.py

@@ -5,30 +5,34 @@ 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")
+    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:
@@ -39,9 +43,9 @@ def predict():
         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"]
-        
+
         # Tokenize input
         inputs = tokenizer(
             code,
@@ -50,21 +54,23 @@ def predict():
             max_length=512,
             return_tensors='pt'
         )
-        
+        print("here")
+
         # 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)