Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
chenpangpang
transformers
Commits
81c8191b
Unverified
Commit
81c8191b
authored
Mar 05, 2024
by
Younes Belkada
Committed by
GitHub
Mar 05, 2024
Browse files
FIX [`Generation`] Fix some issues when running the MaxLength criteria on CPU (#29317)
fix the bitwise or issue
parent
e9476832
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
3 additions
and
3 deletions
+3
-3
src/transformers/generation/stopping_criteria.py
src/transformers/generation/stopping_criteria.py
+3
-3
No files found.
src/transformers/generation/stopping_criteria.py
View file @
81c8191b
...
@@ -73,7 +73,7 @@ class MaxLengthCriteria(StoppingCriteria):
...
@@ -73,7 +73,7 @@ class MaxLengthCriteria(StoppingCriteria):
f
"maximum length (
{
self
.
max_position_embeddings
}
). Depending on the model, you may observe "
f
"maximum length (
{
self
.
max_position_embeddings
}
). Depending on the model, you may observe "
"exceptions, performance degradation, or nothing at all."
"exceptions, performance degradation, or nothing at all."
)
)
return
torch
.
full
((
input_ids
.
shape
[
0
],),
is_done
,
device
=
input_ids
.
device
)
return
torch
.
full
((
input_ids
.
shape
[
0
],),
is_done
,
device
=
input_ids
.
device
,
dtype
=
torch
.
bool
)
class
MaxNewTokensCriteria
(
StoppingCriteria
):
class
MaxNewTokensCriteria
(
StoppingCriteria
):
...
@@ -103,7 +103,7 @@ class MaxNewTokensCriteria(StoppingCriteria):
...
@@ -103,7 +103,7 @@ class MaxNewTokensCriteria(StoppingCriteria):
@
add_start_docstrings
(
STOPPING_CRITERIA_INPUTS_DOCSTRING
)
@
add_start_docstrings
(
STOPPING_CRITERIA_INPUTS_DOCSTRING
)
def
__call__
(
self
,
input_ids
:
torch
.
LongTensor
,
scores
:
torch
.
FloatTensor
,
**
kwargs
)
->
torch
.
BoolTensor
:
def
__call__
(
self
,
input_ids
:
torch
.
LongTensor
,
scores
:
torch
.
FloatTensor
,
**
kwargs
)
->
torch
.
BoolTensor
:
is_done
=
input_ids
.
shape
[
-
1
]
>=
self
.
max_length
is_done
=
input_ids
.
shape
[
-
1
]
>=
self
.
max_length
return
torch
.
full
((
input_ids
.
shape
[
0
],),
is_done
,
device
=
input_ids
.
device
)
return
torch
.
full
((
input_ids
.
shape
[
0
],),
is_done
,
device
=
input_ids
.
device
,
dtype
=
torch
.
bool
)
class
MaxTimeCriteria
(
StoppingCriteria
):
class
MaxTimeCriteria
(
StoppingCriteria
):
...
@@ -126,7 +126,7 @@ class MaxTimeCriteria(StoppingCriteria):
...
@@ -126,7 +126,7 @@ class MaxTimeCriteria(StoppingCriteria):
@
add_start_docstrings
(
STOPPING_CRITERIA_INPUTS_DOCSTRING
)
@
add_start_docstrings
(
STOPPING_CRITERIA_INPUTS_DOCSTRING
)
def
__call__
(
self
,
input_ids
:
torch
.
LongTensor
,
scores
:
torch
.
FloatTensor
,
**
kwargs
)
->
torch
.
BoolTensor
:
def
__call__
(
self
,
input_ids
:
torch
.
LongTensor
,
scores
:
torch
.
FloatTensor
,
**
kwargs
)
->
torch
.
BoolTensor
:
is_done
=
time
.
time
()
-
self
.
initial_timestamp
>
self
.
max_time
is_done
=
time
.
time
()
-
self
.
initial_timestamp
>
self
.
max_time
return
torch
.
full
((
input_ids
.
shape
[
0
],),
is_done
,
device
=
input_ids
.
device
)
return
torch
.
full
((
input_ids
.
shape
[
0
],),
is_done
,
device
=
input_ids
.
device
,
dtype
=
torch
.
bool
)
class
StoppingCriteriaList
(
list
):
class
StoppingCriteriaList
(
list
):
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment