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
OpenDAS
Megatron-LM
Commits
f65a0f88
Commit
f65a0f88
authored
Sep 21, 2021
by
rprenger
Browse files
addressing comments
parent
cb57c380
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
3 additions
and
3 deletions
+3
-3
megatron/text_generation_server.py
megatron/text_generation_server.py
+3
-3
No files found.
megatron/text_generation_server.py
View file @
f65a0f88
...
@@ -23,7 +23,7 @@ from megatron import mpu
...
@@ -23,7 +23,7 @@ from megatron import mpu
from
megatron.text_generation_utils
import
generate
from
megatron.text_generation_utils
import
generate
GENERATE_NUM
=
0
GENERATE_NUM
=
0
sem
=
threading
.
Semaphore
()
lock
=
threading
.
Semaphore
()
class
MegatronGenerate
(
Resource
):
class
MegatronGenerate
(
Resource
):
def
__init__
(
self
,
model
):
def
__init__
(
self
,
model
):
...
@@ -71,10 +71,10 @@ class MegatronGenerate(Resource):
...
@@ -71,10 +71,10 @@ class MegatronGenerate(Resource):
if
not
isinstance
(
add_BOS
,
bool
):
if
not
isinstance
(
add_BOS
,
bool
):
return
"add_BOS must be a boolean value"
return
"add_BOS must be a boolean value"
sem
.
acquire
()
# Need to get lock to keep multiple threads from hitting code
lock
.
acquire
()
# Need to get lock to keep multiple threads from hitting code
MegatronGenerate
.
send_do_generate
()
# Tell other ranks we're doing generate
MegatronGenerate
.
send_do_generate
()
# Tell other ranks we're doing generate
resp_sentences
,
resp_sentences_seg
,
output_logits
,
full_logits
,
tokens
=
generate
(
self
.
model
,
sentences
,
tokens_to_generate
,
all_probs
,
temperature
,
add_BOS
)
resp_sentences
,
resp_sentences_seg
,
output_logits
,
full_logits
,
tokens
=
generate
(
self
.
model
,
sentences
,
tokens_to_generate
,
all_probs
,
temperature
,
add_BOS
)
sem
.
release
()
lock
.
release
()
if
all_probs
:
if
all_probs
:
return
jsonify
({
"sentences"
:
resp_sentences
,
return
jsonify
({
"sentences"
:
resp_sentences
,
...
...
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