Update main.py

main.py

@@ -261,9 +261,9 @@ def predict_vulnerability(code):
             else:
                 outputs = model(**inputs)
         
-        amplified_logits = 3.0 * outputs.logits
+        amplified_logits = 5.0 * outputs.logits
         score = torch.sigmoid(amplified_logits).cpu().item()
-        return round(max(0.0, min(1.0, score)), 4)
+        return 1.0 - round(max(0.0, min(1.0, score)), 4)
         
     except Exception as e:
         print(f"Single prediction error: {e}")
@@ -296,10 +296,10 @@ def predict_vulnerability_batch(codes):
             else:
                 outputs = model(**inputs)
         
-        amplified_logits = 3.0 * outputs.logits
+        amplified_logits = 5.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()]
+        return 1.0 - [round(max(0.0, min(1.0, float(score))), 4) for score in scores.flatten()]
         
     except Exception as e:
         print(f"Batch prediction error: {e}")