refactor: remove comment and create model name

news-copilot-models/main.py

@@ -3,58 +3,48 @@
 from transformers import TextDataset, DataCollatorForLanguageModeling
 from transformers import Trainer, TrainingArguments
 
-# Check device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load tokenizer and model
-tokenizer = GPT2Tokenizer.from_pretrained("news-copilot-viwiki")
-model = GPT2LMHeadModel.from_pretrained("news-copilot-viwiki").to(device)
+model_name = "news-copilot-viwiki"
+
+tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+model = GPT2LMHeadModel.from_pretrained(model_name).to(device)
 
-# Path to your training dataset
 train_file = "./corpus.txt"
 
-# Load dataset
-train_dataset = TextDataset(
-    tokenizer=tokenizer, file_path=train_file, block_size=128  # Adjust as needed
-)
+train_dataset = TextDataset(tokenizer=tokenizer, file_path=train_file, block_size=128)
 
-# Data collator for language modeling
 data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
-# Training arguments
 training_args = TrainingArguments(
     output_dir="./gpt2-finetuned",
     overwrite_output_dir=True,
-    num_train_epochs=1,  # Reduce epochs for faster training
-    per_device_train_batch_size=2,  # Reduce batch size
-    per_device_eval_batch_size=2,  # Reduce batch size
+    num_train_epochs=1,
+    per_device_train_batch_size=2,
+    per_device_eval_batch_size=2,
     logging_dir="./logs",
-    logging_steps=100,  # Reduce logging frequency
+    logging_steps=100,
     evaluation_strategy="steps",
-    eval_steps=500,  # Reduce evaluation frequency
-    save_steps=500,  # Reduce saving frequency
-    warmup_steps=100,  # Reduce warmup steps
+    eval_steps=500,
+    save_steps=500,
+    warmup_steps=100,
     weight_decay=0.01,
     logging_first_step=True,
-    save_total_limit=1,  # Save only one model checkpoint
-    dataloader_num_workers=1,  # Reduce number of workers to 1 for dataloading
+    save_total_limit=1,
+    dataloader_num_workers=1,
 )
 
-# Trainer
 trainer = Trainer(
     model=model,
     args=training_args,
     data_collator=data_collator,
     train_dataset=train_dataset,
 )
 
-# Train the model
 trainer.train()
 
-# Save the fine-tuned model
-trainer.save_model("./gpt2-finetuned")
+trainer.save_model(f"{model_name}-finetuned")
 
-# Save tokenizer
-tokenizer.save_pretrained("./gpt2-finetuned")
+tokenizer.save_pretrained(f"{model_name}-finetuned")
 
-print("Model and tokenizer saved to: ./gpt2-finetuned")
+print("Model and tokenizer saved")