Upload model

config.json

@@ -7,21 +7,6 @@
     "AutoModelForCausalLM": "modeling_gpt.GPTModelForTextGeneration"
   },
   "block_size": 1024,
-  "custom_pipelines": {
-    "text-generation": {
-      "default": {
-        "model": {
-          "pt": "samkeet/GPT_124M-Instruct"
-        }
-      },
-      "impl": "pipeline_gpt.GPT124MTextGenerationPipeline",
-      "pt": [
-        "AutoModelForCausalLM"
-      ],
-      "tf": [],
-      "type": "text"
-    }
-  },
   "model_type": "custom_gpt",
   "n_embd": 768,
   "n_head": 12,

modeling_gpt.py

@@ -284,6 +284,9 @@ class GPTModelForTextGeneration(PreTrainedModel):
             else:
                 next_tokens = torch.argmax(logits, dim=-1, keepdim=True)
 
+            if next_tokens.item() == 50256:
+                break
+
             tokens = torch.cat((tokens, next_tokens), dim=1)
 
         return tokens.flatten()