Update modeling_my_gpt.py

modeling_my_gpt.py

@@ -2,25 +2,24 @@ import torch
 import torch.nn as nn
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import CausalLMOutput
+from transformers.generation import GenerationMixin
+
 
 from .configuration_my_gpt import MyGPTConfig
 from .untrained_model import GPTModel
 
 import os
-import sys
-
-curr_dir = os.getcwd()
-parent_dir = os.path.dirname(curr_dir)
 
-sys.path.insert(0, parent_dir)
 
 
-class MyGPTForCausalLM(PreTrainedModel):
+class MyGPTForCausalLM(PreTrainedModel, GenerationMixin):
     config_class = MyGPTConfig
+    main_input_name = "input_ids"
 
     def __init__(self, config):
         super().__init__(config)
 
+
         # Import your original GPTModel
         self.model = GPTModel({
             "vocab_size": config.vocab_size,