Changes epochs to 1 and batch size to 8 and disabled gradient checkpoint

app.py

@@ -7,7 +7,6 @@ import traceback
 
 dir_path = os.path.abspath('./')
 os.environ["HF_HOME"] = dir_path
-start_training = st.button("Train Model")
 
 
 def tokenize_function(examples):
@@ -19,69 +18,70 @@ def tokenize_function(examples):
     return tokenized_inputs
 
 
-if start_training:
-    st.write("Getting model and dataset ...")
-    # Load the dataset
-    dataset = load_dataset("viber1/indian-law-dataset", cache_dir=dir_path)
+st.write("Getting model and dataset ...")
+# Load the dataset
+dataset = load_dataset("viber1/indian-law-dataset", cache_dir=dir_path)
 
-    # Update this path based on where the tokenizer files are actually stored
-    tokenizer = AutoTokenizer.from_pretrained('gpt2')
-    tokenizer.pad_token = tokenizer.eos_token
-    # Load the model
-    model = AutoModelForCausalLM.from_pretrained('gpt2')
-    model.gradient_checkpointing_enable()
+# Update this path based on where the tokenizer files are actually stored
+tokenizer = AutoTokenizer.from_pretrained('gpt2')
+tokenizer.pad_token = tokenizer.eos_token
+# Load the model
+model = AutoModelForCausalLM.from_pretrained('gpt2')
+# model.gradient_checkpointing_enable()
 
-    st.write("Training setup ...")
-    # Apply the tokenizer to the dataset
-    tokenized_dataset = dataset.map(tokenize_function, batched=True)
+st.write("Training setup ...")
+# Apply the tokenizer to the dataset
+tokenized_dataset = dataset.map(tokenize_function, batched=True)
 
-    # Split the dataset manually into train and validation sets
-    split_dataset = tokenized_dataset["train"].train_test_split(test_size=0.1)
+# Split the dataset manually into train and validation sets
+split_dataset = tokenized_dataset["train"].train_test_split(test_size=0.1)
 
-    # Convert the dataset to PyTorch tensors
-    train_dataset = split_dataset["train"].with_format("torch")
-    eval_dataset = split_dataset["test"].with_format("torch")
+# Convert the dataset to PyTorch tensors
+train_dataset = split_dataset["train"].with_format("torch")
+eval_dataset = split_dataset["test"].with_format("torch")
 
-    # Create data loaders
-    # reduce batch size 8 to 1
-    train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True, pin_memory=True)
-    eval_dataloader = DataLoader(eval_dataset, batch_size=1, pin_memory=True)
+# Create data loaders
+# reduce batch size 8 to 1
+train_dataloader = DataLoader(train_dataset, batch_size=8, shuffle=True, pin_memory=True)
+eval_dataloader = DataLoader(eval_dataset, batch_size=8, pin_memory=True)
 
-    # Define training arguments
-    training_args = TrainingArguments(
-        output_dir="./results",
-        eval_strategy="epoch",
-        learning_rate=2e-5,
-        per_device_train_batch_size=1,
-        per_device_eval_batch_size=1,
-        num_train_epochs=3,
-        weight_decay=0.01,
-        fp16=True,  # Enable mixed precision
-        # save_total_limit=2,
-        logging_dir='./logs',  # Set logging directory
-        logging_steps=10,  # Log more frequently
-        gradient_checkpointing=True,  # Enable gradient checkpointing
-        gradient_accumulation_steps=8  # Accumulate gradients over 8
-    )
+# Define training arguments
+training_args = TrainingArguments(
+    output_dir="{dir_path}/results",
+    eval_strategy="epoch",
+    learning_rate=2e-5,
+    per_device_train_batch_size=8,
+    per_device_eval_batch_size=8,
+    num_train_epochs=1,
+    weight_decay=0.01,
+    fp16=True,  # Enable mixed precision
+    # save_total_limit=2,
+    logging_dir='{dir_path}/logs',  # Set logging directory
+    logging_steps=5,  # Log more frequently
+    # gradient_checkpointing=True,  # Enable gradient checkpointing
+    # gradient_accumulation_steps=8  # Accumulate gradients over 8
+)
 
-    st.write("Training Started .....")
+st.write("Training Started .....")
 
-    # Create the Trainer
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=train_dataset,
-        eval_dataset=eval_dataset,
-    )
-
-    try:
-        trainer.train()
-    except Exception as e:
-        st.write(f"Error: {e}")
-        traceback.print_exc()
-        st.write("some error")
+# Create the Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+)
 
+try:
+    trainer.train()
     # Evaluate the model
     st.write("Training Done ...")
+    model.save_pretrained(f"{dir_path}\\trained-gpt2")
+    tokenizer.save_pretrained(f"{dir_path}\\trained-gpt2")
+    st.write("Evaluating Model ...")
     results = trainer.evaluate()
     st.write(results)
+except Exception as e:
+    st.write(f"Error: {e}")
+    traceback.print_exc()
+    st.write("some error")