Update main.py

main.py

@@ -261,11 +261,8 @@ def predict_readability(code):
             else:
                 outputs = model(**inputs)
         
-        if hasattr(outputs, 'logits'):
-            score = torch.sigmoid(outputs.logits).cpu().item()
-        else:
-            score = torch.sigmoid(outputs[0]).cpu().item()
-        
+        amplified_logits = 4.0 * outputs.logits
+        score = torch.sigmoid(amplified_logits).cpu().item()
         return round(max(0.0, min(1.0, score)), 4)
         
     except Exception as e:
@@ -299,11 +296,9 @@ def predict_readability_batch(codes):
             else:
                 outputs = model(**inputs)
         
-        if hasattr(outputs, 'logits'):
-            scores = torch.sigmoid(outputs.logits).cpu().numpy()
-        else:
-            scores = torch.sigmoid(outputs[0]).cpu().numpy()
-        
+        amplified_logits = 4.0 * outputs.logits
+        scores = torch.sigmoid(amplified_logits).cpu().numpy()
+
         return [round(max(0.0, min(1.0, float(score))), 4) for score in scores.flatten()]
         
     except Exception as e: