update text server to support compute logprobs (NVIDIA#7733)

* update text server to support compute logprobs * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix typo --------- Signed-off-by: Zhilin Wang <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
anmolgupt · Oct 24, 2023 · 265a0a6 · 265a0a6
1 parent 609af14
commit 265a0a6
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/nemo/collections/nlp/modules/common/text_generation_server.py b/nemo/collections/nlp/modules/common/text_generation_server.py
@@ -45,6 +45,7 @@
         "repetition_penalty",
         "min_tokens_to_generate",
         "end_strings",
+        "compute_logprob",
     ]
 )
 
@@ -165,6 +166,12 @@ def put(self):
             if neighbors < 0:
                 return "num of neighbors must be an integer no less than 0"
 
+        compute_logprob = False
+        if "compute_logprob" in request.get_json():
+            compute_logprob = request.get_json()["compute_logprob"]
+            if not isinstance(compute_logprob, bool):
+                return "compute_logprob must be a boolean value"
+
         with lock:  # Need to get lock to keep multiple threads from hitting code
             MegatronGenerate.send_do_generate()  # Tell other ranks we're doing generate
             extra = {}
@@ -192,6 +199,7 @@ def put(self):
                 repetition_penalty,
                 end_strings=end_strings,
                 min_tokens_to_generate=min_tokens_to_generate,
+                compute_logprob=compute_logprob,
                 **extra,
             )
             for k in output: