Update train.py

train.py

@@ -1,3 +1,8 @@
+import os
+os.environ["HF_HOME"] = "./hf_cache"
+os.environ["TRANSFORMERS_CACHE"] = "./hf_cache/transformers"
+os.environ["HF_DATASETS_CACHE"] = "./hf_cache/datasets"
+
 from datasets import load_dataset
 from transformers import AutoTokenizer, DataCollatorForLanguageModeling, Trainer, TrainingArguments, AutoModelForMaskedLM
 
@@ -37,4 +42,4 @@ trainer = Trainer(
 )
 
 # Start training
-trainer.train()
+trainer.train()