Reduced the number of parameters

app.py

@@ -13,8 +13,13 @@ dataset = load_dataset("mwitiderrick/swahili")
 # Print dataset columns for verification
 print(f"Dataset columns: {dataset['train'].column_names}")
 
+# Select a subset of the dataset (e.g., first 100,000 rows)
+subset_size = 100000  # Adjust the size as needed
+subset_dataset = dataset["train"].select(range(min(subset_size, len(dataset["train"]))))
+print(f"Using a subset of {len(subset_dataset)} rows for training.")
+
 # Initialize the tokenizer and model
-model_name = "gpt2"  # Use a smaller variant of GPT-2 for efficiency
+model_name = "gpt2"  # Use GPT-2 for text generation
 tokenizer = GPT2Tokenizer.from_pretrained(model_name)
 model = GPT2LMHeadModel.from_pretrained(model_name)
 
@@ -24,21 +29,21 @@ tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
 
 # Preprocess the dataset
 def preprocess_function(examples):
+    # Tokenize and format the dataset
     encodings = tokenizer(
-        examples['text'], 
+        examples['text'],  # Use 'text' column from your dataset
         truncation=True,
-        padding='max_length', 
+        padding='max_length',  # Ensure consistent length
         max_length=512
     )
-    encodings['labels'] = encodings['input_ids']
+    encodings['labels'] = encodings['input_ids']  # Use input_ids directly as labels
     return encodings
 
 # Tokenize the dataset
 try:
-    tokenized_datasets = dataset.map(
+    tokenized_datasets = subset_dataset.map(
         preprocess_function, 
-        batched=True,
-        batch_size=1000  # Adjust batch size for efficiency
+        batched=True
     )
 except Exception as e:
     print(f"Error during tokenization: {e}")
@@ -46,21 +51,20 @@ except Exception as e:
 # Define training arguments
 training_args = TrainingArguments(
     output_dir='./results',
-    per_device_train_batch_size=2,  # Lowered batch size to prevent OOM errors
+    per_device_train_batch_size=4,
     num_train_epochs=1,
     logging_dir='./logs',
-    logging_steps=500,
-    evaluation_strategy="steps",
-    save_steps=5000,  # Save checkpoints more frequently
-    save_total_limit=2,
-    gradient_accumulation_steps=8,  # Accumulate gradients to simulate larger batch size
+    logging_steps=500,  # Log every 500 steps
+    evaluation_strategy="steps",  # Use evaluation strategy
+    save_steps=10_000,  # Save checkpoint every 10,000 steps
+    save_total_limit=2,  # Keep only the last 2 checkpoints
 )
 
 # Define Trainer
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=tokenized_datasets["train"],
+    train_dataset=tokenized_datasets,
     tokenizer=tokenizer,
 )