From ba0252b7f9c3048e14e2e3d215580ce601eeaf4c Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Fri, 25 Dec 2020 14:20:35 +0100 Subject: [PATCH] Cleaning up conversation tests. --- src/transformers/pipelines/conversational.py | 12 +++- tests/test_pipelines_common.py | 63 ++++++++++++++++++++ tests/test_pipelines_conversational.py | 51 +++++++++++++++- tests/test_pipelines_summarization.py | 35 +---------- 4 files changed, 124 insertions(+), 37 deletions(-) diff --git a/src/transformers/pipelines/conversational.py b/src/transformers/pipelines/conversational.py index a9ee449f..4bc4c213 100644 --- a/src/transformers/pipelines/conversational.py +++ b/src/transformers/pipelines/conversational.py @@ -2,8 +2,8 @@ import uuid from typing import List, Optional, Union from ..file_utils import add_end_docstrings, is_tf_available, is_torch_available -from ..utils import logging from ..tokenization_utils import TruncationStrategy +from ..utils import logging from .base import PIPELINE_INIT_ARGS, Pipeline @@ -48,12 +48,18 @@ class Conversation: conversation.add_user_input("Is it good?") """ - def __init__(self, text: str = None, conversation_id: uuid.UUID = None): + def __init__( + self, text: str = None, conversation_id: uuid.UUID = None, past_user_inputs=None, generated_responses=None + ): if not conversation_id: conversation_id = uuid.uuid4() self.uuid: uuid.UUID = conversation_id + if past_user_inputs is None: + past_user_inputs = past_user_inputs self.past_user_inputs: List[str] = [] - self.generated_responses: List[str] = [] + if generated_responses is None: + generated_responses = [] + self.generated_responses: List[str] = generated_responses self.history: List[int] = [] self.new_user_input: Optional[str] = text diff --git a/tests/test_pipelines_common.py b/tests/test_pipelines_common.py index c8a66053..951aa96a 100644 --- a/tests/test_pipelines_common.py +++ b/tests/test_pipelines_common.py @@ -18,9 +18,14 @@ from unittest import mock from transformers import is_tf_available, is_torch_available, pipeline from transformers.pipelines import Pipeline from transformers.testing_utils import _run_slow_tests, is_pipeline_test, require_tf, require_torch, slow +from transformers.tokenization_utils import TruncationStrategy from transformers.tokenization_utils_base import to_py_obj +if is_torch_available(): + import torch + + VALID_INPUTS = ["A simple string", ["list of strings"]] @@ -242,3 +247,61 @@ class MonoInputPipelineCommonMixin(CustomInputPipelineCommonMixin): self.assertIn(key, result) self.assertRaises(Exception, nlp, self.invalid_inputs) + + +class DummyTok: + pad_token_id = 0 + eos_token_id = 0 + + def __init__(self, **kwargs): + for name, v in kwargs.items(): + setattr(self, name, v) + + def __call__(self, inputs, **kwargs): + if kwargs.get("return_tensors", "") == "pt": + return self.encode_pt(inputs, **kwargs) + else: + return self.encode_list(inputs, **kwargs) + + def encode_list(self, inputs, **kwargs): + unwrap = False + if isinstance(inputs, str): + unwrap = True + inputs = [inputs] + + assert isinstance(inputs, list) + input_ids = [self.encode(input_) for input_ in inputs] + + if unwrap: + input_ids = input_ids[0] + + return {"input_ids": input_ids} + + def encode_pt(self, inputs, **kwargs): + if isinstance(inputs, str): + input_ids = torch.LongTensor(self.encode(inputs)).unsqueeze(0) + else: + input_ids = self._pad([self.encode(input_) for input_ in inputs]) + return self.finalize_pt(input_ids, **kwargs) + + def finalize_pt(self, input_ids, **kwargs): + if kwargs.get("truncation", TruncationStrategy.DO_NOT_TRUNCATE) == TruncationStrategy.ONLY_FIRST: + input_ids = input_ids[:, : self.model_max_length] + attention_mask = torch.zeros_like(input_ids).long() + 1 + return {"input_ids": input_ids, "attention_mask": attention_mask} + + def _pad(self, inputs): + return torch.nn.utils.rnn.pad_sequence( + [torch.LongTensor(input_) for input_ in inputs], + padding_value=self.pad_token_id, + ).transpose(1, 0) + + def pad(self, inputs, **kwargs): + input_ids = self._pad(inputs["input_ids"]) + return self.finalize_pt(input_ids, **kwargs) + + def encode(self, input_): + return list(input_.encode("utf-8")) + + def decode(self, sequence, **kwargs): + return "D" * len(sequence) diff --git a/tests/test_pipelines_conversational.py b/tests/test_pipelines_conversational.py index 16703b11..10be49d0 100644 --- a/tests/test_pipelines_conversational.py +++ b/tests/test_pipelines_conversational.py @@ -15,14 +15,63 @@ import unittest from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, Conversation, ConversationalPipeline, pipeline +from transformers.models.gpt2 import GPT2Config, GPT2LMHeadModel from transformers.testing_utils import require_torch, slow, torch_device -from .test_pipelines_common import MonoInputPipelineCommonMixin +from .test_pipelines_common import DummyTok, MonoInputPipelineCommonMixin DEFAULT_DEVICE_NUM = -1 if torch_device == "cpu" else 0 +class SimpleConversationPipelineTests(unittest.TestCase): + @require_torch + def test_integration_torch_conversation(self): + # When + config = GPT2Config( + vocab_size=257, n_ctx=64, n_embd=64, n_layer=1, n_head=8, eos_token_id=0, bos_token_id=0, pad_token_id=0 + ) + model = GPT2LMHeadModel(config) + tokenizer = DummyTok() + nlp = pipeline(task="conversational", device=DEFAULT_DEVICE_NUM, model=model, tokenizer=tokenizer) + conversation_1 = Conversation("Going to the movies tonight - any suggestions?") + conversation_2 = Conversation("What's the last book you have read?") + # Then + self.assertEqual(len(conversation_1.past_user_inputs), 0) + self.assertEqual(len(conversation_2.past_user_inputs), 0) + # When + result = nlp([conversation_1, conversation_2], do_sample=False, max_length=1000) + # Then + self.assertEqual(result, [conversation_1, conversation_2]) + self.assertEqual( + result, + [ + Conversation( + None, + past_user_inputs=["Going to the movies tonight - any suggestions?"], + generated_responses=["D"], + ), + Conversation( + None, past_user_inputs=["What's the last book you have read?"], generated_responses=["D"] + ), + ], + ) + + # When + conversation_2.add_user_input("Why do you recommend it?") + result = nlp(conversation_2, do_sample=False, max_length=1000) + # Then + self.assertEqual(result, conversation_2) + self.assertEqual( + result, + Conversation( + None, + past_user_inputs=["What's the last book you have read?", "Why do you recommend it?"], + generated_responses=["D", "D"], + ), + ) + + class ConversationalPipelineTests(MonoInputPipelineCommonMixin, unittest.TestCase): pipeline_task = "conversational" small_models = [] # Models tested without the @slow decorator diff --git a/tests/test_pipelines_summarization.py b/tests/test_pipelines_summarization.py index 78d312bb..813a0324 100644 --- a/tests/test_pipelines_summarization.py +++ b/tests/test_pipelines_summarization.py @@ -14,47 +14,16 @@ import unittest -from transformers import is_torch_available, pipeline +from transformers import pipeline from transformers.models.bart import BartConfig, BartForConditionalGeneration from transformers.testing_utils import require_torch, slow, torch_device from transformers.tokenization_utils import TruncationStrategy -from .test_pipelines_common import MonoInputPipelineCommonMixin +from .test_pipelines_common import DummyTok, MonoInputPipelineCommonMixin DEFAULT_DEVICE_NUM = -1 if torch_device == "cpu" else 0 -if is_torch_available(): - import torch - - -class DummyTok: - pad_token_id = 0 - - def __init__(self, **kwargs): - for name, v in kwargs.items(): - setattr(self, name, v) - - def __call__(self, inputs, **kwargs): - if isinstance(inputs, str): - input_ids = self.encode(inputs).unsqueeze(0) - else: - input_ids = torch.nn.utils.rnn.pad_sequence( - [self.encode(input_) for input_ in inputs], - padding_value=self.pad_token_id, - ) - - if kwargs.get("truncation", TruncationStrategy.DO_NOT_TRUNCATE) == TruncationStrategy.ONLY_FIRST: - input_ids = input_ids[:, : self.model_max_length] - attention_mask = torch.zeros_like(input_ids).long() + 1 - return {"input_ids": input_ids, "attention_mask": attention_mask} - - def encode(self, input_): - return torch.LongTensor(list(input_.encode("utf-8"))) - - def decode(self, sequence, **kwargs): - return "D" * len(sequence) - class SimpleSummarizationPipelineTests(unittest.TestCase): @require_torch