Create train_and_save_model.py

train_and_save_model.py

@@ -0,0 +1,89 @@
+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
+for epoch in range(num_epochs):
+    for inputs, targets in data_loader:
+        # Forward pass
+        outputs = model(inputs)
+        loss = criterion(outputs, targets)
+
+        # Backward pass and optimization
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+    
+    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!")