Update train_and_save_model.py

train_and_save_model.py

@@ -1,99 +1,54 @@
-import torch
-import torch.nn as nn
-import torch.optim as optim
-from torch.utils.data import DataLoader, Dataset
-import json
-import os
-
-# Step 1: Define Your Dataset Class
-class CustomDataset(Dataset):
-    def __init__(self, texts, labels):
-        self.texts = texts
-        self.labels = labels
-
-    def __len__(self):
-        return len(self.texts)
-
-    def __getitem__(self, idx):
-        return self.texts[idx], self.labels[idx]
-
-# Step 2: Define Your Model Class
-class LSTMModel(nn.Module):
-    def __init__(self, input_size, hidden_size, output_size):
-        super(LSTMModel, self).__init__()
-        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
-        self.fc = nn.Linear(hidden_size, output_size)
-
-    def forward(self, x):
-        lstm_out, _ = self.lstm(x)
-        out = self.fc(lstm_out[:, -1, :])  # Get the last time step output
-        return out
-
-# Step 3: Initialize Hyperparameters and Model
-input_size = 100  # Example input size (e.g., embedding size)
-hidden_size = 64  # Number of LSTM units
-output_size = 10  # Number of output classes
-num_epochs = 5
-learning_rate = 0.001
-
-# Initialize the model
-model = LSTMModel(input_size, hidden_size, output_size)
-
-# Step 4: Set Up Loss and Optimizer
-criterion = nn.CrossEntropyLoss()
-optimizer = optim.Adam(model.parameters(), lr=learning_rate)
-
-# Step 5: Sample Data (You would replace this with your actual data)
-# Here, we create random data for demonstration purposes
-texts = torch.randn(100, 10, input_size)  # 100 samples, sequence length of 10
-labels = torch.randint(0, output_size, (100,))  # 100 random labels
-
-# Create a DataLoader
-dataset = CustomDataset(texts, labels)
-data_loader = DataLoader(dataset, batch_size=16, shuffle=True)
-
-# Step 6: Training Loop
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)  # Move model to the correct device
-
-for epoch in range(num_epochs):
-    for inputs, targets in data_loader:
-        # Move data to the same device as the model
-        inputs, targets = inputs.to(device), targets.to(device)
-
-        # Forward pass
-        outputs = model(inputs)
-        loss = criterion(outputs, targets)
-
-        # Backward pass and optimization
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-
-        # Print loss for every batch
-        print(f'Epoch [{epoch+1}/{num_epochs}], Batch Loss: {loss.item():.4f}')
-    
-    # Print epoch summary
-    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
-
-# Step 7: Save the Model
-model_save_path = "path/to/save/model_directory"  # Change this to your desired path
-os.makedirs(model_save_path, exist_ok=True)  # Create the directory if it doesn't exist
-
-# Save the model weights
-torch.save(model.state_dict(), os.path.join(model_save_path, "pytorch_model.bin"))
-
-# Step 8: Create and Save the Configuration File
-config = {
-    "input_size": input_size,
-    "hidden_size": hidden_size,
-    "output_size": output_size,
-    "num_layers": 1,  # Add more parameters as needed
-    "dropout": 0.2
-}
-
-# Save the configuration to a JSON file
-with open(os.path.join(model_save_path, "config.json"), "w") as f:
-    json.dump(config, f)
-
-print("Model and configuration saved successfully!")
+from datasets import load_dataset
+from transformers import AutoAdapterModel, AutoTokenizer, Trainer, TrainingArguments
+
+# Load datasets
+dataset_pentesting = load_dataset("canstralian/pentesting-ai")
+dataset_redpajama = load_dataset("togethercomputer/RedPajama-Data-1T")
+
+# Tokenizer
+tokenizer = AutoTokenizer.from_pretrained("canstralian/rabbitredeux")
+
+def tokenize_function(examples):
+    return tokenizer(examples['text'], padding="max_length", truncation=True)
+
+# Tokenize datasets
+tokenized_dataset_pentesting = dataset_pentesting.map(tokenize_function, batched=True)
+tokenized_dataset_redpajama = dataset_redpajama.map(tokenize_function, batched=True)
+
+# Prepare datasets
+train_dataset_pentesting = tokenized_dataset_pentesting["train"]
+validation_dataset_pentesting = tokenized_dataset_pentesting["validation"]
+
+# Load model and adapter
+model = AutoAdapterModel.from_pretrained("canstralian/rabbitredeux")
+model.load_adapter("Canstralian/RabbitRedux", set_active=True)
+
+# Training arguments
+training_args = TrainingArguments(
+    output_dir="./results",
+    num_train_epochs=3,
+    per_device_train_batch_size=8,
+    per_device_eval_batch_size=8,
+    warmup_steps=500,
+    weight_decay=0.01,
+    logging_dir="./logs",
+    logging_steps=10,
+    evaluation_strategy="epoch",
+)
+
+# Trainer setup
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset_pentesting,
+    eval_dataset=validation_dataset_pentesting,
+)
+
+# Training
+trainer.train()
+
+# Evaluate model
+trainer.evaluate()
+
+# Save the fine-tuned model
+model.save_pretrained("./fine_tuned_model")