Commit f65a0f88 authored by rprenger's avatar rprenger
Browse files

addressing comments

parent cb57c380
...@@ -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,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment