Updating README and splitting training logic

README.md

@@ -48,10 +48,9 @@ This project implements a transformer-based language model using PyTorch. The mo
    git clone https://github.com/yourusername/transformer-model-training.git
    cd transformer-model-training
    ```
-
-2. Install the required packages:
+2. To train the model, run the training script:
    ```bash
-   pip install -r requirements.txt
+   python train.py
    ```
 
 ## Usage

app.py

@@ -9,7 +9,8 @@ def load_model():
     config = GPTConfig()
     model = GPT(config)
     try:
-        model.load_state_dict(torch.load('trained_model_quantized.pt'))
+        # Load the model with map_location to handle CPU-only environments
+        model.load_state_dict(torch.load('trained_model_quantized.pt', map_location=torch.device('cpu')), strict=False)
         model.eval()  # Set the model to evaluation mode
         st.success("Model loaded successfully!")
     except Exception as e:
@@ -24,7 +25,7 @@ def load_tokenizer():
 def generate_text(model, tokenizer, input_text, length, num_sequences):
     # Encode the input text
     input_ids = tokenizer.encode(input_text)
-    input_tensor = torch.tensor(input_ids).unsqueeze(0)  # Add batch dimension
+    input_tensor = torch.tensor(input_ids).unsqueeze(0)  # Add batch dimension (shape: [1, T])
 
     generated_sequences = []
     for _ in range(num_sequences):
@@ -35,7 +36,10 @@ def generate_text(model, tokenizer, input_text, length, num_sequences):
                 next_token_logits = logits[:, -1, :]  # Get the last token's logits
                 next_token_probs = torch.softmax(next_token_logits, dim=-1)
                 next_token = torch.multinomial(next_token_probs, num_samples=1)  # Sample from the distribution
-                input_tensor = torch.cat((input_tensor, next_token.unsqueeze(0)), dim=1)  # Append the new token
+                
+                # Ensure the next_token has the correct shape for concatenation
+                next_token = next_token.view(1, -1)  # Reshape to [1, 1] if necessary
+                input_tensor = torch.cat((input_tensor, next_token), dim=1)  # Append the new token
 
         # Decode the generated tokens
         generated_sequences.append(tokenizer.decode(input_tensor[0].tolist()))
@@ -51,8 +55,8 @@ length = st.slider("Predict Additional Text of Length", 1, 50, 10)
 num_sequences = st.slider("Number of Sequences to Generate", 1, 5, 1)
 
 if st.button("Generate"):
-    model = load_model()
-    tokenizer = load_tokenizer()
+    model = load_model()  # Load the model for inference
+    tokenizer = load_tokenizer()  # Load the tokenizer
     st.write("Generating text...")
     generated_texts = generate_text(model, tokenizer, input_text, length, num_sequences)
     st.write("Text generation complete.")

train.py

@@ -0,0 +1,106 @@
+import os
+import time
+import torch
+from transformer import GPT, GPTConfig, DataLoaderLite  # Import your model and data loader
+
+# Initialize the model and data loader
+config = GPTConfig()
+model = GPT(config)
+train_loader = DataLoaderLite(B=4, T=1024)
+
+# Define the optimizer
+optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
+
+# Function to load the most recent checkpoint
+def load_latest_checkpoint(model):
+    checkpoint_file = 'checkpoint.pt'
+    if not os.path.exists(checkpoint_file):
+        return 0  # No checkpoint found, start from epoch 0
+
+    print(f'Loading checkpoint from {checkpoint_file}')
+    checkpoint = torch.load(checkpoint_file)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    return checkpoint['epoch']
+
+# Load the latest checkpoint if available
+start_epoch = load_latest_checkpoint(model)
+
+# Training loop
+num_epochs = 78
+
+# Start time tracking
+start_time = time.time()
+
+for epoch in range(start_epoch, num_epochs):  # Start from the loaded epoch
+    epoch_loss = 0.0  # Initialize epoch loss
+    num_steps = 0  # Initialize step counter for the epoch
+    last_loss = None  # Variable to store the last loss
+
+    # Calculate total steps for the progress bar
+    total_steps = len(train_loader.tokens) // (train_loader.B * train_loader.T)
+
+    # Use tqdm to create a progress bar
+    with tqdm(total=total_steps, desc=f'Epoch {epoch + 1}/{num_epochs}') as pbar:
+        for step in range(total_steps):  # Iterate over the number of steps
+            x, y = train_loader.next_batch()
+            x, y = x.to(device), y.to(device)
+            optimizer.zero_grad()
+            logits, loss = model(x, y)
+            loss.backward()
+            optimizer.step()
+            
+            epoch_loss += loss.item()  # Accumulate loss
+            num_steps += 1  # Increment step counter
+            last_loss = loss.item()  # Store the last loss
+            pbar.update(1)  # Update progress bar
+
+            # Check if the loss is below the threshold
+            if last_loss < 0.099999:
+                print(f'Loss below threshold: {last_loss:.6f}')  # Print loss before breaking
+                break  # Exit the loop if the loss condition is met
+
+    # Print the loss at the end of the epoch
+    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {last_loss:.6f}')
+
+    # Check if the loss condition was met to break out of the epoch loop
+    if last_loss < 0.099999:
+        print(f'Early stopping at epoch {epoch + 1} due to loss condition met.')
+        break  # Exit the epoch loop if the loss condition is met
+
+    # Checkpointing: Save the model and the current epoch after each epoch
+    checkpoint_path = 'checkpoint.pt'  # Save to a single checkpoint file
+    torch.save({
+        'epoch': epoch + 1,  # Save the current epoch number
+        'model_state_dict': model.state_dict(),  # Save the model state
+    }, checkpoint_path)
+    print(f'Checkpoint saved to {checkpoint_path}')
+
+# End time tracking
+end_time = time.time()
+training_duration = end_time - start_time
+
+# Convert training duration to minutes and seconds
+minutes = int(training_duration // 60)
+seconds = int(training_duration % 60)
+
+# Print the total training time in minute:second format
+print(f'Total training time: {minutes} minutes and {seconds} seconds')
+
+# After training your model, apply quantization and save it with compression
+def save_model_with_quantization(model, file_path):
+    # Switch model to evaluation mode
+    model.eval()
+    
+    # Apply dynamic quantization
+    quantized_model = torch.quantization.quantize_dynamic(
+        model,  # the model to be quantized
+        {nn.Linear},  # layers to quantize
+        dtype=torch.qint8  # quantization type
+    )
+    
+    # Save the quantized model with compression
+    torch.save(quantized_model.state_dict(), file_path, _use_new_zipfile_serialization=True)
+    print(f'Model saved to {file_path} with quantization and compression.')
+
+# Call this function after training your model
+save_model_with_quantization(model, 'trained_model_quantized.pt')

transformer.py

@@ -233,121 +233,3 @@ class DataLoaderLite:
         if self.current_position + (B * T + 1) > len(self.tokens):
             self.current_position = 0
         return x, y
-
-# Initialize the data loader with batch size 4 and sequence length 1024
-train_loader = DataLoaderLite(B=4, T=1024)
-
-# Initialize the model
-model = GPT(GPTConfig())
-model.to(device)
-
-# Print number of model parameters
-model.print_num_parameters()
-
-# Define the optimizer
-optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
-
-# Function to load the most recent checkpoint
-def load_latest_checkpoint(model):
-    # Find the checkpoint file
-    checkpoint_file = 'checkpoint.pt'
-    if not os.path.exists(checkpoint_file):
-        return 0  # No checkpoint found, start from epoch 0
-
-    print(f'Loading checkpoint from {checkpoint_file}')
-    
-    # Load the model state and epoch number
-    checkpoint = torch.load(checkpoint_file)
-    
-    # Ensure the checkpoint contains the expected keys
-    if 'model_state_dict' not in checkpoint or 'epoch' not in checkpoint:
-        raise KeyError("Checkpoint does not contain required keys.")
-
-    model.load_state_dict(checkpoint['model_state_dict'])
-    
-    # Return the epoch number
-    return checkpoint['epoch']
-
-# Load the latest checkpoint if available
-start_epoch = load_latest_checkpoint(model)
-
-# NEW CODE: Training loop until loss is less than 0.099999
-loss = float('inf')  # Initialize loss to a large value
-num_epochs = 78  # Set the number of epochs to 78
-
-# Start time tracking
-start_time = time.time()
-
-for epoch in range(start_epoch, num_epochs):  # Start from the loaded epoch
-    epoch_loss = 0.0  # Initialize epoch loss
-    num_steps = 0  # Initialize step counter for the epoch
-    last_loss = None  # Variable to store the last loss
-
-    # Calculate total steps for the progress bar
-    total_steps = len(train_loader.tokens) // (train_loader.B * train_loader.T)
-
-    # Use tqdm to create a progress bar
-    with tqdm(total=total_steps, desc=f'Epoch {epoch + 1}/{num_epochs}') as pbar:
-        for step in range(total_steps):  # Iterate over the number of steps
-            x, y = train_loader.next_batch()
-            x, y = x.to(device), y.to(device)
-            optimizer.zero_grad()
-            logits, loss = model(x, y)
-            loss.backward()
-            optimizer.step()
-            
-            epoch_loss += loss.item()  # Accumulate loss
-            num_steps += 1  # Increment step counter
-            last_loss = loss.item()  # Store the last loss
-            pbar.update(1)  # Update progress bar
-
-            # Check if the loss is below the threshold
-            if last_loss < 0.099999:
-                print(f'Loss below threshold: {last_loss:.6f}')  # Print loss before breaking
-                break  # Exit the loop if the loss condition is met
-
-    # Print the loss at the end of the epoch
-    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {last_loss:.6f}')
-
-    # Check if the loss condition was met to break out of the epoch loop
-    if last_loss < 0.099999:
-        print(f'Early stopping at epoch {epoch + 1} due to loss condition met.')
-        break  # Exit the epoch loop if the loss condition is met
-
-    # Checkpointing: Save the model and the current epoch after each epoch
-    checkpoint_path = 'checkpoint.pt'  # Save to a single checkpoint file
-    torch.save({
-        'epoch': epoch + 1,  # Save the current epoch number
-        'model_state_dict': model.state_dict(),  # Save the model state
-    }, checkpoint_path)
-    print(f'Checkpoint saved to {checkpoint_path}')
-
-# End time tracking
-end_time = time.time()
-training_duration = end_time - start_time
-
-# Convert training duration to minutes and seconds
-minutes = int(training_duration // 60)
-seconds = int(training_duration % 60)
-
-# Print the total training time in minute:second format
-print(f'Total training time: {minutes} minutes and {seconds} seconds')
-
-# After training your model, apply quantization and save it with compression
-def save_model_with_quantization(model, file_path):
-    # Switch model to evaluation mode
-    model.eval()
-    
-    # Apply dynamic quantization
-    quantized_model = torch.quantization.quantize_dynamic(
-        model,  # the model to be quantized
-        {nn.Linear},  # layers to quantize
-        dtype=torch.qint8  # quantization type
-    )
-    
-    # Save the quantized model with compression
-    torch.save(quantized_model.state_dict(), file_path, _use_new_zipfile_serialization=True)
-    print(f'Model saved to {file_path} with quantization and compression.')
-
-# Call this function after training your model
-save_model_with_quantization(model, 'trained_model_quantized.pt')