added comments and more explicit training loop

qatransformer2.py

@@ -1,10 +1,12 @@
 from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForQuestionAnswering, TrainingArguments, Trainer, DefaultDataCollator
+from transformers.optimization import AdamW
+from transformers.data.data_collator import default_data_collator
 
 squad = load_dataset("squad", split="train[:5000]")
 squad = squad.train_test_split(test_size=0.2)
 
-# preprocess
+
 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 
 def preprocess_function(examples):
@@ -58,13 +60,16 @@ def preprocess_function(examples):
     inputs["end_positions"] = end_positions
     return inputs
 
+# Define the model
+model = AutoModelForQuestionAnswering.from_pretrained("bert-base-uncased")
 
+# Define the optimization algorithm
+optimizer = AdamW(model.parameters(), lr=2e-5)
 
-# train
-train_dataset = squad["train"].map(preprocess_function, batched=True)
-eval_dataset = squad["test"].map(preprocess_function, batched=True)
+# Define the loss function
+loss_fn = default_data_collator
 
-model = AutoModelForQuestionAnswering.from_pretrained("bert-base-uncased")
+# Define the training arguments
 training_args = TrainingArguments(
     output_dir="question-answering",
     evaluation_strategy="epoch",
@@ -76,16 +81,20 @@ training_args = TrainingArguments(
     push_to_hub=True,
 )
 
+# Define the trainer
 trainer = Trainer(
     model=model,
     args=training_args,
     train_dataset=train_dataset,
     eval_dataset=eval_dataset,
     tokenizer=tokenizer,
-    data_collator=data_collator,
+    data_collator=loss_fn,
+    optimizer=optimizer,
 )
 
+# Train the model
 trainer.train()
 
 
+
 #evaluation - todo