Update train_script.py

Added loading pretrained model

train_script.py

@@ -4,7 +4,7 @@ import os
 import torch
 import torch.nn as nn
 
-from transformers import GPT2TokenizerFast, GPT2LMHeadModel
+from transformers import GPT2TokenizerFast, GPT2LMHeadModel, AutoModelForCausalLM
 from transformers import DataCollatorWithPadding, GPT2Config, DataCollatorForLanguageModeling
 from transformers import Trainer, TrainingArguments, RobertaTokenizerFast
 
@@ -58,7 +58,7 @@ dataset = dataset.with_format("torch")
 tokenizer = RobertaTokenizerFast.from_pretrained(ENCODER_MODEL_NAME, max_len=TOKENIZER_MAX_LEN)
 collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
 
-
+# train from scratch
 config = GPT2Config(
     vocab_size=len(tokenizer),
     n_positions=TOKENIZER_MAX_LEN,
@@ -71,6 +71,11 @@ config = GPT2Config(
 
 model = ConditionalGPT2LMHeadModel(config)
 
+# alternatively, load a pre-trained model
+# commit_hash = '0ba58478f467056fe33003d7d91644ecede695a7'
+# model = AutoModelForCausalLM.from_pretrained("entropy/roberta_zinc_decoder",
+                                              # trust_remote_code=True, revision=commit_hash)
+
 # change trainer args as needed
 args = TrainingArguments(
     output_dir=TRAINER_SAVE_DIR,