Skip to content

Commit

Permalink
Re-add unit test
Browse files Browse the repository at this point in the history
  • Loading branch information
anmarques committed Nov 8, 2023
1 parent 28919b1 commit 21c6f0d
Showing 1 changed file with 23 additions and 0 deletions.
23 changes: 23 additions & 0 deletions tests/deepsparse/transformers/pipelines/test_text_generation.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,29 @@ def test_token_generation_non_deterministic(pipeline, prompt):
assert len(set(text_outputs)) == 3


def test_pipeline_for_ppl_eval(self, ):
pipeline = self.get_pipeline(
task="text-generation",
model_path=self.model_stub,
sequence_length=self.sequence_length,
prompt_sequence_length=1,
)
inputs = dict(
prompt=self.prompt,
output_scores=True,
return_input_tokens=True,
fixed_sequences_length=True,
include_prompt_logits=True,
max_length=1,
)
predictions = pipeline(**inputs)
assert hasattr(predictions, "generations")
assert hasattr(predictions.generations[0], "score")
assert hasattr(predictions.generations[0], "input_tokens")
assert "input_ids" in predictions.generations[0].input_tokens
assert "attention_mask" in predictions.generations[0].input_tokens


def test_streaming_mode_returns_generator(pipeline, prompt):
response_generator = pipeline(prompt, streaming=True)
assert inspect.isgenerator(
Expand Down

0 comments on commit 21c6f0d

Please sign in to comment.