Update script.py

script.py

@@ -2,79 +2,54 @@ import subprocess
 import sys
 
 print("Installing required packages and loading model...")
-process = subprocess.Popen(
-    [sys.executable, "-m", "pip", "install", "-q", "transformers", "accelerate", "peft", "torch"],
-    stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
-)
 processx = subprocess.Popen(
     [sys.executable, "-m", "pip", "install", "-q", "-U", "bitsandbytes",],
     stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
 )
-for line in process.stdout:
-    print(line, end='')
-process.wait()
 for line in processx.stdout:
     print(line, end='')
 processx.wait()
         
+process = subprocess.Popen(
+    [sys.executable, "-m", "pip", "install", "-q", "transformers", "accelerate", "peft", "torch"],
+    stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
+)
+for line in process.stdout:
+    print(line, end='')
+process.wait()
+
         
-import torch
-import random
 import ast
 import re
-import threading
 from peft import PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
-# Global variables to store the model, tokenizer, and installation state
-_model = None
-_tokenizer = None
-_model_lock = threading.Lock()
-_initialized = False  # Flag to track initialization
+bnb_config = BitsAndBytesConfig(load_in_8bit=True)
 
-def initialize():
-    """Install required packages and load the model/tokenizer only once."""
-    global _model, _tokenizer, _initialized
-    if _initialized:
-        return  # Prevent duplicate initialization
-    
-    with _model_lock:
-        if _initialized:
-            return  # Double-check inside the lock to prevent race conditions
-        
-        # Enable quantization to reduce memory usage
-        bnb_config = BitsAndBytesConfig(load_in_8bit=True)
-        
-        # Load tokenizer
-        _tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B-Instruct")
-        
-        # Load model with quantization
-        base_model = AutoModelForCausalLM.from_pretrained(
-            "Qwen/Qwen2.5-Coder-7B-Instruct",
-            quantization_config=bnb_config,
-            device_map="auto",
-        )
-        
-        # Load the fine-tuned model
-        _model = PeftModel.from_pretrained(
-            base_model,
-            "SushantGautam/vulnerability_ativ0.1",
-            device_map="auto",
-        )
-        
-        _initialized = True
+# Load tokenizer
+_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B-Instruct")
+
+# Load model with quantization
+base_model = AutoModelForCausalLM.from_pretrained(
+    "Qwen/Qwen2.5-Coder-7B-Instruct",
+    quantization_config=bnb_config,
+    device_map="auto",
+)
+
+# Load the fine-tuned model
+_model = PeftModel.from_pretrained(
+    base_model,
+    "SushantGautam/vulnerability_ativ0.1",
+    device_map="auto",
+)
 
-def load_model():
-    """Ensure model is initialized before returning it."""
-    initialize()
-    return _model, _tokenizer
 
 def extract_dict(text):
     match = re.search(r"```python\n(.*?)\n```", text, re.DOTALL)
-    return ast.literal_eval(match.group(1)) if match else None
+    return ast.literal_eval(match.group(1)) if match else {}
 
 def generate(prompt):
-    model, tokenizer = load_model()
+    model, tokenizer = _model, _tokenizer
     messages = [
         {"role": "system", "content": "You are a cybersecurity expert specializing in CWE vulnerabilities in codes. Your responses must be accompanied by a python JSON."},
         {"role": "user", "content": prompt},
@@ -100,7 +75,7 @@ def generate(prompt):
     try:
         response_formatted = extract_dict(response)
     except:
-        response_formatted = None
+        response_formatted = {}
     
     return {"Generated Answer": response, "Extracted Dict": response_formatted}
 
@@ -110,3 +85,5 @@ print("💪🏆🎉 Pong! Model and tokenizer loaded successfully. Use generate(
 
 
 # prompt = "Here's a properly secured code snippet:\n\ndef add_label options, f, attr\n    label_size = options.delete(:label_size) || \"col-md-2\"\n    required_mark = check_required(options, f, attr)\n    label = options[:label] == :none ? '' : options.delete(:label)\n    label ||= ((clazz = f.object.class).respond_to?(:gettext_translation_for_attribute_name) &&\n        s_(clazz.gettext_translation_for_attribute_name attr)) if f\n    label = label.present? ? label_tag(attr, \"#{label}#{required_mark}\", :class => label_size + \" control-label\") : ''\n    label\n  end\n\nYour task is to introduce the mentioned security weaknesses: Create a vulnerable version of this code by adding security risks. Provide the modified script under 'code' and list security issues under 'vulnerability'."
+
+# 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)