Create script.py

script.py

@@ -0,0 +1,87 @@
+import torch
+import random
+import ast
+import re
+import subprocess
+import sys
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+# Global variables to store the model and tokenizer (to avoid multiple loads)
+
+_model = None
+_tokenizer = None
+
+
+def load_model():
+    global _model, _tokenizer
+    if _model is None or _tokenizer is None:
+        print("Installing required packages...")
+        process = subprocess.Popen(
+            [sys.executable, "-m", "pip", "install", "bitsandbytes", "transformers", "accelerate", "peft", "torch"],
+            stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
+        )
+        for line in process.stdout: print(line, end='')
+        process.wait()
+
+        print("Loading model and tokenizer...")
+        # 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",
+        )
+    else:
+        print("Model and tokenizer already loaded.")
+    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
+
+
+def generate(prompt):
+    model, tokenizer = load_model()
+    
+    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},
+    ]
+    
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    generated_ids = model.generate(
+        **model_inputs,
+        max_new_tokens=4000,
+    )
+    generated_ids = [
+        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+    ]
+    
+    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    
+    try:
+        response_formatted = extract_dict(response)
+    except:
+        response_formatted = "XXX"
+    
+    return {"Generated Answer": response, "Extracted Dict": response_formatted}