(return_log_prob)
| 171 | |
| 172 | @pytest.mark.parametrize("return_log_prob", [True, False]) |
| 173 | def test_token_streaming(return_log_prob): |
| 174 | source = ["آ", "ت", "ز", "م", "و", "ن"] |
| 175 | translator = _get_transliterator() |
| 176 | |
| 177 | expected_result = translator.translate_batch([source], return_scores=True)[0] |
| 178 | |
| 179 | step_results = translator.generate_tokens(source, return_log_prob=return_log_prob) |
| 180 | assert inspect.isgenerator(step_results) |
| 181 | |
| 182 | tokens = [] |
| 183 | cum_log_probs = 0 |
| 184 | |
| 185 | for step_result in step_results: |
| 186 | assert isinstance(step_result, ctranslate2.GenerationStepResult) |
| 187 | |
| 188 | tokens.append(step_result.token) |
| 189 | |
| 190 | if return_log_prob: |
| 191 | cum_log_probs += step_result.log_prob |
| 192 | else: |
| 193 | assert step_result.log_prob is None |
| 194 | |
| 195 | assert tokens == expected_result.hypotheses[0] + ["</s>"] |
| 196 | |
| 197 | if return_log_prob: |
| 198 | assert cum_log_probs / len(tokens) == pytest.approx( |
| 199 | expected_result.scores[0], abs=1e-5 |
| 200 | ) |
| 201 | |
| 202 | |
| 203 | def test_token_streaming_exception(): |
nothing calls this directly
no test coverage detected