Merge branch 'develop' of https://github.com/PaddlePaddle/PaddleNLP into update_predict_new

Author: DesmonDay

.markdownlint.yaml

@@ -248,16 +248,19 @@ def _preprocess(self, source):
     def _infer(self, inputs):
         raise NotImplementedError
 
-    def _postprocess(self, predictions):
+    def _postprocess(self, predictions, return_tokens=False):
         decoded_predictions = self.tokenizer.batch_decode(
             predictions, skip_special_tokens=True, clean_up_tokenization_spaces=False
         )
-        return decoded_predictions
+        if return_tokens:
+            return decoded_predictions, predictions
+        else:
+            return decoded_predictions
 
-    def predict(self, input_texts: str | list[str]):
+    def predict(self, input_texts: str | list[str], return_tokens=False):
         tokenized_source = self._preprocess(input_texts)
         predictions = self._infer(tokenized_source)
-        decoded_predictions = self._postprocess(predictions)
+        decoded_predictions = self._postprocess(predictions, return_tokens=return_tokens)
         return decoded_predictions
 
 
@@ -470,13 +473,16 @@ def __init__(self, config: PredictorArgument, tokenizer: PretrainedTokenizer):
             )
             self.generation_config = None
 
-    def _postprocess(self, predictions):
+    def _postprocess(self, predictions, return_tokens=False):
         if paddle.distributed.get_rank() == 0:
             tokens: np.ndarray = load_real_time_tokens()
             decoded_predictions = self.tokenizer.batch_decode(
                 tokens.tolist(), skip_special_tokens=True, clean_up_tokenization_spaces=False
             )
-            return decoded_predictions
+            if return_tokens:
+                return decoded_predictions, tokens.tolist()
+            else:
+                return decoded_predictions
         else:
             return None
 
@@ -1034,7 +1040,7 @@ def _infer(self, inputs: dict[str, paddle.Tensor]):
         )
 
     @paddle.no_grad()
-    def predict(self, input_texts: str | list[str]):
+    def predict(self, input_texts: str | list[str], return_tokens=False):
         self._preprocess(input_texts)
 
         result_queue = mp.Queue()
@@ -1055,9 +1061,15 @@ def predict(self, input_texts: str | list[str]):
             self.used_list[i] = []
 
         outputs = []
+        output_tokens = []
         while len(outputs) < self.batch_size:
-            outputs.append(result_queue.get(timeout=1)[-1])
-        return outputs
+            result = result_queue.get(timeout=1)
+            outputs.append(result[-1])
+            output_tokens.append(result[-2])
+        if return_tokens:
+            return outputs, output_tokens
+        else:
+            return outputs
 
 
 class StaticBlockInferencePredictor(BlockInferencePredictorMixin, BasePredictor):
@@ -1180,7 +1192,7 @@ def _share_data(self):
     def _infer(self):
         self.predictor.run()
 
-    def predict(self, input_texts: str | list[str]):
+    def predict(self, input_texts: str | list[str], return_tokens=False):
 
         s_time = time.time()
         self._preprocess(input_texts)
@@ -1213,9 +1225,15 @@ def predict(self, input_texts: str | list[str]):
             self.used_list[i] = []
 
         outputs = []
+        output_tokens = []
         while len(outputs) < self.batch_size:
-            outputs.append(result_queue.get(timeout=1)[-1])
-        return outputs
+            result = result_queue.get(timeout=1)
+            outputs.append(result[-1])
+            output_tokens.append(result[-2])
+        if return_tokens:
+            return outputs, output_tokens
+        else:
+            return outputs
 
     def _preprocess(self, source):
         BlockInferencePredictorMixin._preprocess(self, source)
@@ -1681,8 +1699,8 @@ def benchmark(predictor, predictor_args, model_args):
     output_tokens = 0
     for _ in range(test_time):
         for bs, batch_source_text in enumerate(batch_benchmark_texts):
-            outputs = predictor.predict(batch_source_text)
-            output_tokens += sum([len(output) for output in outputs])
+            outputs, batch_tokens = predictor.predict(batch_source_text, return_tokens=True)
+            output_tokens += sum([len(tokens) for tokens in batch_tokens])
     end = time.perf_counter()
     print("Avg Elapse time is: ", (end - start) / test_time)
     print("Output tokens is: ", output_tokens)

llm/predict/predictor.py

@@ -797,7 +797,7 @@ def read_res(model_name_or_path: str, tensor_queue: mp.Queue, result_queue: mp.Q
             break
     output = np.concatenate(outputs, axis=1).tolist()
     seqs = tokenizer.batch_decode(output, skip_special_tokens=True, clean_up_tokenization_spaces=False)
-    for i, seq in enumerate(seqs):
-        result_queue.put([i, seq])
+    for i, (out, seq) in enumerate(zip(output, seqs)):
+        result_queue.put([i, out, seq])
 
     logger.info("Finish read result message")

llm/utils/utils.py

@@ -94,10 +94,11 @@ def get_datasets_weights_and_num_samples(data_prefix, train_val_test_num_samples
     # Add 0.5% (the 1.005 factor) so in case the bleding dataset does
     # not uniformly distribute the number of samples, we still have
     # samples left to feed to the network.
+    # (NOTE, yujun06): This is a workaround to avoid issues with indexing in the blending dataset. Therefore, we need to add 20 samples to each dataset.
     datasets_train_valid_test_num_samples = []
     for weight in weights:
         datasets_train_valid_test_num_samples.append(
-            [int(math.ceil(val * weight * 1.005)) for val in train_val_test_num_samples]
+            [int(math.ceil(val * weight * 1.005)) + 20 for val in train_val_test_num_samples]
         )
 
     return prefixes, weights, datasets_train_valid_test_num_samples

paddlenlp/data/causal_dataset.py

@@ -1211,6 +1211,8 @@ def sample(
                 probs = TopPProcess(probs, top_p, min_tokens_to_keep)
             if paddle.device.is_compiled_with_custom_device("gcu"):
                 probs = paddle.cast(probs, "float32")
+            if paddle.device.is_compiled_with_xpu():
+                probs = paddle.cast(probs, "float32")
 
             # multinomial already support fp16 and bf16 currently, fix issue: https://github.com/PaddlePaddle/Paddle/issues/51852
             next_tokens = paddle.multinomial(probs)