Update script.py

script.py

@@ -1,11 +1,12 @@
 import subprocess
 import sys
 import json
+from retrying import retry
 
 print("🔥 Installing required packages...")
         
 process = subprocess.Popen(
-    [sys.executable, "-m", "pip", "install", "-q", "transformers", "bitsandbytes", "accelerate", "peft", "torch"],
+    [sys.executable, "-m", "pip", "install", "-q", "transformers", "retrying", "bitsandbytes", "accelerate", "peft", "torch"],
     stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
 )
 for line in process.stdout:
@@ -41,6 +42,7 @@ def extract_dict(text):
     match = re.search(r"```python\n(.*?)\n```", text, re.DOTALL)
     return ast.literal_eval(match.group(1)) if match else {}
 
+@retry(stop_max_attempt_number=5)
 def generate(prompt, max_new_tokens=4000):
     model, tokenizer = _model, _tokenizer
     messages = [
@@ -67,10 +69,14 @@ def generate(prompt, max_new_tokens=4000):
     
     try:
         response_formatted = extract_dict(response)
+        assert "code" in response_formatted or "vulnerability" in response_formatted
+        if "vulnerability" in response_formatted:
+            assert isinstance(response_formatted["vulnerability"], list)
     except:
-        response_formatted = {}
-    
-    return {"Generated Answer": response, "Extracted Dict": response_formatted}
+        print('⚠️ Failed to extract dict properly from response. Retrying...')
+        raise ValueError(f"❌ Failed to extract dict properly from response after five tries: {response}")
+    return {"raw": response, "extracted": {response_formatted}}
+
 print("🎉 Model and tokenizer loaded successfully. Now, making sure the model responses: Ping!")
 generate("Ping!")
 print('💪🏆🎉 Pong! Model responded successfully. Use generate("<prompt>") to call the model.')
@@ -80,7 +86,6 @@ print('💪🏆🎉 Pong! Model responded successfully. Use generate("<prompt>")
 # Use: globals().get("_x") or exec(requests.get("https://huggingface.co/SushantGautam/vulnerability_ativ0.1/raw/main/script.py").text) or globals().__setitem__("_x", 1)
 
 
-
 def vulnerability_injection(secure_code: str, CWE_vulnerability=None):
     if CWE_vulnerability is None:
          ## introduce any vulnerability
@@ -94,7 +99,7 @@ def vulnerability_injection(secure_code: str, CWE_vulnerability=None):
             """
     else:
      ## introduce specific vulnerability
-     prompt = f"""
+        prompt = f"""
             Modify the following secure code to introduce a {CWE_vulnerability} vulnerability.
 
             Secure Code:
@@ -104,7 +109,8 @@ def vulnerability_injection(secure_code: str, CWE_vulnerability=None):
             Return JSON with keys: 'code' (modified vulnerable code) and 'vulnerability' (list of CWE if vulnerabilities introduced else empty).
 
             """
-    return generate(prompt)
+    return generate(prompt)    
+
 
 
 
@@ -115,7 +121,6 @@ def vulnerability_detection(input_code: str):
         Code:
         {input_code}
 
-
         Identify potential vulnerabilities that could be introduced. Your task is to assess this script for potential vulnerabilities that could be introduced.
         Return JSON with key: 'vulnerability' (list of detected CWE vulnerabilities or an empty list if none found).
         """