Skip to content

Commit

Permalink
Fix and re-enable ConversationalPipeline tests (#26907)
Browse files Browse the repository at this point in the history
* Fix and re-enable conversationalpipeline tests

* Fix the batch test so the change only applies to conversational pipeline
  • Loading branch information
Rocketknight1 authored Oct 19, 2023
1 parent 734dd96 commit bdbcd5d
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 7 deletions.
8 changes: 4 additions & 4 deletions tests/pipelines/test_pipelines_conversational.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,14 +77,14 @@ def get_test_pipeline(self, model, tokenizer, processor):

def run_pipeline_test(self, conversation_agent, _):
# Simple
outputs = conversation_agent(Conversation("Hi there!"))
outputs = conversation_agent(Conversation("Hi there!"), max_new_tokens=20)
self.assertEqual(
outputs,
Conversation([{"role": "user", "content": "Hi there!"}, {"role": "assistant", "content": ANY(str)}]),
)

# Single list
outputs = conversation_agent([Conversation("Hi there!")])
outputs = conversation_agent([Conversation("Hi there!")], max_new_tokens=20)
self.assertEqual(
outputs,
Conversation([{"role": "user", "content": "Hi there!"}, {"role": "assistant", "content": ANY(str)}]),
Expand All @@ -96,7 +96,7 @@ def run_pipeline_test(self, conversation_agent, _):
self.assertEqual(len(conversation_1), 1)
self.assertEqual(len(conversation_2), 1)

outputs = conversation_agent([conversation_1, conversation_2])
outputs = conversation_agent([conversation_1, conversation_2], max_new_tokens=20)
self.assertEqual(outputs, [conversation_1, conversation_2])
self.assertEqual(
outputs,
Expand All @@ -118,7 +118,7 @@ def run_pipeline_test(self, conversation_agent, _):

# One conversation with history
conversation_2.add_message({"role": "user", "content": "Why do you recommend it?"})
outputs = conversation_agent(conversation_2)
outputs = conversation_agent(conversation_2, max_new_tokens=20)
self.assertEqual(outputs, conversation_2)
self.assertEqual(
outputs,
Expand Down
9 changes: 6 additions & 3 deletions tests/test_pipeline_mixin.py
Original file line number Diff line number Diff line change
Expand Up @@ -312,8 +312,12 @@ def data(n):
yield copy.deepcopy(random.choice(examples))

out = []
for item in pipeline(data(10), batch_size=4):
out.append(item)
if task == "conversational":
for item in pipeline(data(10), batch_size=4, max_new_tokens=20):
out.append(item)
else:
for item in pipeline(data(10), batch_size=4):
out.append(item)
self.assertEqual(len(out), 10)

run_batch_test(pipeline, examples)
Expand All @@ -327,7 +331,6 @@ def test_pipeline_automatic_speech_recognition(self):
self.run_task_tests(task="automatic-speech-recognition")

@is_pipeline_test
@unittest.skip("Conversational tests are currently broken for several models, will fix ASAP - Matt")
def test_pipeline_conversational(self):
self.run_task_tests(task="conversational")

Expand Down

0 comments on commit bdbcd5d

Please sign in to comment.