Create main.py

main.py

@@ -0,0 +1,16 @@
+from datasets import load_dataset
+from transformers import GPT2LMHeadModel
+
+dataset = load_dataset("Mireu-Lab/UNSW-NB15", name="en-US")
+model = GPT2LMHeadModel.from_pretrained("gpt2")
+
+# 训练模型
+trainer = Trainer(
+    model=model,
+    train_dataset=train_data,
+)
+
+trainer.train()
+
+# 评估模型
+trainer.evaluate()