File tree 2 files changed +6
-2
lines changed 2 files changed +6
-2
lines changed Original file line number Diff line number Diff line change @@ -31,7 +31,7 @@ def __init__(
31
31
use_kv_cache : bool = False ,
32
32
):
33
33
device = "cuda" if torch .cuda .is_available () else "cpu"
34
- super ().__init__ (device = device )
34
+ super ().__init__ (device = device , pretrained = "gpt2" )
35
35
self ._model = model
36
36
self ._tokenizer = tokenizer
37
37
self ._device = torch .device (device )
@@ -47,6 +47,10 @@ def eot_token_id(self):
47
47
return self ._tokenizer .eot_id
48
48
return self ._tokenizer .eos_id
49
49
50
+ @property
51
+ def prefix_token_id (self ):
52
+ return self ._tokenizer .eos_id
53
+
50
54
@property
51
55
def max_length (self ):
52
56
return self ._max_seq_length
Original file line number Diff line number Diff line change @@ -15,7 +15,7 @@ pip install --no-use-pep517 "git+https://github.com/pytorch/ao.git@${TORCHAO_VER
15
15
16
16
# Install lm-eval for Model Evaluation with lm-evalution-harness
17
17
# Install tiktoken for tokenizer
18
- pip install lm_eval==0.4.2
18
+ pip install lm_eval==0.4.5
19
19
pip install tiktoken blobfile
20
20
21
21
# Call the install helper for further setup
You can’t perform that action at this time.
0 commit comments