Update main.py

main.py

@@ -187,7 +187,7 @@ def predict_batch():
                         batch_scores[local_idx] = 0.0
                 
                 for score in batch_scores:
-                    results.append({"score": 1.0 - score})
+                    results.append({"score": round(1.0 - score,4)})
                 
                 cleanup_gpu_memory()
         
@@ -261,7 +261,7 @@ def predict_vulnerability(code):
             else:
                 outputs = model(**inputs)
         
-        amplified_logits = 4.0 * outputs.logits
+        amplified_logits = 2.0 * outputs.logits
         score = torch.sigmoid(amplified_logits).cpu().item()
         return round(max(0.0, min(1.0, score)), 4)
         
@@ -296,7 +296,7 @@ def predict_vulnerability_batch(codes):
             else:
                 outputs = model(**inputs)
         
-        amplified_logits = 4.0 * outputs.logits
+        amplified_logits = 2.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()]