Merge pull request #4 from harubaru/main

Add initial GPT-J support
corolla-johnson · Sep 27, 2021 · a25c72d · a25c72d
2 parents aba081c + 3a52860
commit a25c72d
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/mkultra/inference.py b/mkultra/inference.py
@@ -1,10 +1,11 @@
-from transformers import GPT2LMHeadModel, GPTNeoForCausalLM, TextGenerationPipeline
+from transformers import GPT2LMHeadModel, GPTNeoForCausalLM, GPTJForCausalLM, TextGenerationPipeline
 from mkultra.soft_prompt import SoftPrompt
 import torch
 
 EXTRA_ALLOWED_MODELS = [
     "GPT2SoftPromptLM",
-    "GPTNeoSoftPromptLM"
+    "GPTNeoSoftPromptLM",
+    "GPTJSoftPromptLM"
     ]
 
 for model in EXTRA_ALLOWED_MODELS:
@@ -85,5 +86,9 @@ def __init__(self, config):
         super().__init__(config)
 
 class GPTNeoSoftPromptLM(GPTSoftPromptMixin, GPTNeoForCausalLM):
+    def __init__(self, config):
+        super().__init__(config)
+
+class GPTJSoftPromptLM(GPTSoftPromptMixin, GPTJForCausalLM):
     def __init__(self, config):
         super().__init__(config)
diff --git a/mkultra/tuning.py b/mkultra/tuning.py
@@ -1,4 +1,4 @@
-from transformers import GPT2LMHeadModel, GPTNeoForCausalLM
+from transformers import GPT2LMHeadModel, GPTNeoForCausalLM, GPTJForCausalLM
 from mkultra.soft_prompt import SoftPrompt
 import torch
 import torch.nn as nn
@@ -119,5 +119,9 @@ def __init__(self, config):
         super().__init__(config)
 
 class GPTNeoPromptTuningLM(GPTPromptTuningMixin, GPTNeoForCausalLM):
+    def __init__(self, config):
+        super().__init__(config)
+
+class GPTJPromptTuningLM(GPTPromptTuningMixin, GPTJForCausalLM):
     def __init__(self, config):
         super().__init__(config)