fix: gptj engine bug and add dependency

quick fix gptj engine bug and add accelerate dependency
tarunchy · Apr 14, 2023 · 52f4b44 · 52f4b44
1 parent 668f61b
commit 52f4b44
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -59,6 +59,7 @@ dependencies = [
     "openai >= 0.27.0",
     "pydantic >= 1.10.0",
     "rouge-score >= 0.1.2",
+    "accelerate",
 ]
 
 [project.scripts]

diff --git a/src/xturing/engines/gptj_engine.py b/src/xturing/engines/gptj_engine.py
@@ -46,10 +46,8 @@ def __init__(self, weights_path: Optional[Union[str, Path]] = None):
             "philschmid/gpt-j-6B-fp16-sharded", load_in_8bit=True, device_map=device_map
         )
 
-        tokenizer = AutoTokenizer.from_pretrained(
-            "philschmid/gpt-j-6B-fp16-sharded"
-        )
-        tokenizer.pad_token = self.tokenizer.eos_token
+        tokenizer = AutoTokenizer.from_pretrained("philschmid/gpt-j-6B-fp16-sharded")
+        tokenizer.pad_token = tokenizer.eos_token
         super().__init__(
             weights_path=weights_path, model=model, tokenizer=tokenizer, load_8bit=True
         )