# Path: prompt_tuning/data/utils.py
# def identity(x: Any, *args, **kwargs) -> Any:  # pylint: disable=unused-argument
# def task_mapping(tasks: Sequence[str],
#                  aliases: Optional[Mapping[str, str]]) -> Dict[str, int]:
# def remove_prefix(s: str, prefix: str) -> str:
# def remove_suffix(string: str, suffix: str) -> str:


from absl.testing import absltest
from prompt_tuning.data import utils
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for utils."""



class UtilsTest(absltest.TestCase):

  def test_task_mapping(self):
    tasks = ['a', 'alias_a', 'alias_b', 'b', 'c', 'alias_b2']
    aliases = {
        'alias_a': 'a',
        'alias_b': 'b',
        'alias_b2': 'b'
    }
    gold_tasks = {
        'a': 0,
        'b': 1,
        'c': 2,
        'alias_a': 0,
        'alias_b': 1,
        'alias_b2': 1,
    }
    task_index = utils.task_mapping(tasks, aliases)
    self.assertEqual(task_index, gold_tasks)

  def test_task_mapping_raises_error(self):
    with self.assertRaises(ValueError):

# Path: prompt_tuning/data/utils.py
# def identity(x: Any, *args, **kwargs) -> Any:  # pylint: disable=unused-argument
# def task_mapping(tasks: Sequence[str],
#                  aliases: Optional[Mapping[str, str]]) -> Dict[str, int]:
# def remove_prefix(s: str, prefix: str) -> str:
# def remove_suffix(string: str, suffix: str) -> str:


# Path: prompt_tuning/data/utils_test.py
from absl.testing import absltest
from prompt_tuning.data import utils

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for utils."""



class UtilsTest(absltest.TestCase):

  def test_task_mapping(self):
    tasks = ['a', 'alias_a', 'alias_b', 'b', 'c', 'alias_b2']
    aliases = {
        'alias_a': 'a',
        'alias_b': 'b',
        'alias_b2': 'b'
    }
    gold_tasks = {
        'a': 0,
        'b': 1,
        'c': 2,
        'alias_a': 0,
        'alias_b': 1,
        'alias_b2': 1,
    }
    task_index = utils.task_mapping(tasks, aliases)
    self.assertEqual(task_index, gold_tasks)

  def test_task_mapping_raises_error(self):
    with self.assertRaises(ValueError):

      utils.task_mapping([], {'alias': 'missing_task'})