Update app.py

app.py

@@ -33,9 +33,14 @@ model, tokenizer = FastLanguageModel.from_pretrained(
 )
 print("Model and tokenizer loaded successfully.")
 
+# Wrap the model in DataParallel to use all GPUs
+if torch.cuda.device_count() > 1:
+    print(f"Using {torch.cuda.device_count()} GPUs!")
+    model = torch.nn.DataParallel(model)
+
 print("Configuring PEFT model...")
 model = FastLanguageModel.get_peft_model(
-    model,
+    model.module if isinstance(model, torch.nn.DataParallel) else model,
     r=16,
     target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
     lora_alpha=16,
@@ -51,24 +56,18 @@ print("PEFT model configured.")
 # Updated alpaca_prompt for different types
 alpaca_prompt = {
     "learning_from": """Below is a CVE definition.
-
 ### CVE definition:
 {}
-
 ### detail CVE:
 {}""",
     "definition": """Below is a definition about software vulnerability. Explain it.
-
 ### Definition:
 {}
-
 ### Explanation:
 {}""",
     "code_vulnerability": """Below is a code snippet. Identify the line of code that is vulnerable and describe the type of software vulnerability.
-
 ### Code Snippet:
 {}
-
 ### Vulnerability solution:
 {}"""
 }
@@ -111,7 +110,7 @@ print("Formatting function applied.")
 
 print("Initializing trainer...")
 trainer = SFTTrainer(
-    model=model,
+    model=model.module if isinstance(model, torch.nn.DataParallel) else model,
     tokenizer=tokenizer,
     train_dataset=dataset,
     dataset_text_field="text",
@@ -145,11 +144,16 @@ num += 1
 uploads_models = f"cybersentinal-3.0"
 
 print("Saving the trained model...")
-model.save_pretrained_merged("model", tokenizer, save_method="merged_16bit")
+model.module.save_pretrained_merged("model", tokenizer, save_method="merged_16bit") if isinstance(model, torch.nn.DataParallel) else model.save_pretrained_merged("model", tokenizer, save_method="merged_16bit")
 print("Model saved successfully.")
 
 print("Pushing the model to the hub...")
-model.push_to_hub_merged(
+model.module.push_to_hub_merged(
+    uploads_models,
+    tokenizer,
+    save_method="merged_16bit",
+    token=hf_token
+) if isinstance(model, torch.nn.DataParallel) else model.push_to_hub_merged(
     uploads_models,
     tokenizer,
     save_method="merged_16bit",