Merge pull request caikit#255 from gkumbhat/add_data_limitation

Add data limitation

caikit_nlp/config/config.yml

@@ -30,5 +30,10 @@ unload_tgis_prompt_artifacts: false
 master_addr: localhost
 master_port: 29550
 
+training_data_limit:
+ # Configuration for PeftPromptTuning module
+ 6655831b-960a-4dc5-8df4-867026e2cd41:
+ add_model_name_here: 10000
+
 runtime:
  library: caikit_nlp

caikit_nlp/modules/text_generation/peft_prompt_tuning.py

@@ -41,6 +41,7 @@
 import transformers
 
 # First Party
+from caikit import get_config
 from caikit.core.data_model import DataStream
 from caikit.core.exceptions import error_handler
 from caikit.core.modules import ModuleBase, ModuleConfig, ModuleSaver, module
@@ -366,6 +367,20 @@ def train(
  verbalizer,
  )
 
+ # Check if data is within limit allowed for this module and model
+ max_num_examples = (
+ get_config()
+ .training_data_limit.get(cls.MODULE_ID, {})
+ .get(base_model_name, -1)
+ )
+
+ if max_num_examples > 0:
+ error.value_check(
+ "<NLP77627434E>",
+ len(train_stream) <= max_num_examples,
+ "Number of examples larger than maximum number of examples allowed for this model",
+ )
+
  # Coerce the passed model into a resource; if we have one, this is a noop
  # TODO: When splitting up this mono-module, use the configured resource
  # type of the concrete class to bootstrap

pyproject.toml

@@ -14,7 +14,7 @@ classifiers=[
  "License :: OSI Approved :: Apache Software License"
 ]
 dependencies = [
- "caikit[runtime-grpc,runtime-http]>=0.22.0,<0.23.0",
+ "caikit[runtime-grpc,runtime-http]>=0.23.2,<0.25.0",
  "caikit-tgis-backend>=0.1.17,<0.2.0",
  # TODO: loosen dependencies
  "accelerate>=0.22.0",
@@ -32,7 +32,7 @@ dependencies = [
  # which broke caikit-nlp build. peft hasn't released newer version yet, so to get
  # the build fix, we pulling peft from main branch commit. In future, we will pull PEFT from
  # pypi
- "peft@git+https://github.com/huggingface/peft.git#8c17d556a8fe9522e10d73d7bd3fad46a6ecae14"
+ "peft@git+https://github.com/huggingface/peft.git@8c17d556a8fe9522e10d73d7bd3fad46a6ecae14"
 ]
 
 [tool.setuptools.packages.find]

tests/modules/text_generation/test_peft_prompt_tuning.py

@@ -33,6 +33,7 @@
  seq2seq_lm_dummy_model,
  seq2seq_lm_train_kwargs,
  set_cpu_device,
+ temp_config,
 )
 import caikit_nlp
 
@@ -399,3 +400,108 @@ def test_run_exponential_decay_len_penatly_object(causal_lm_dummy_model):
  exponential_decay_length_penalty=penalty,
  )
  assert isinstance(pred, GeneratedTextResult)
+
+
+def test_train_with_data_validation_raises(causal_lm_train_kwargs, set_cpu_device):
+ """Check if we are able to throw error for when number of examples are more than configured limit"""
+ patch_kwargs = {
+ "num_epochs": 1,
+ "verbalizer": "Tweet text : {{input}} Label : ",
+ "train_stream": caikit.core.data_model.DataStream.from_iterable(
+ [
+ ClassificationTrainRecord(
+ text="@foo what a cute dog!", labels=["no complaint"]
+ ),
+ ClassificationTrainRecord(
+ text="@bar this is the worst idea ever.", labels=["complaint"]
+ ),
+ ]
+ ),
+ "torch_dtype": torch.bfloat16,
+ "device": "cpu",
+ }
+ causal_lm_train_kwargs.update(patch_kwargs)
+
+ model_name = causal_lm_train_kwargs["base_model"]._model_name
+ module = caikit_nlp.modules.text_generation.PeftPromptTuning
+ with temp_config(training_data_limit={module.MODULE_ID: {model_name: 1}}):
+ with pytest.raises(ValueError):
+ module.train(**causal_lm_train_kwargs)
+
+
+def test_train_with_data_validation_success(causal_lm_train_kwargs, set_cpu_device):
+ """Check if we are able to train successfully if training data is within limits"""
+ patch_kwargs = {
+ "num_epochs": 1,
+ "verbalizer": "Tweet text : {{input}} Label : ",
+ "train_stream": caikit.core.data_model.DataStream.from_iterable(
+ [
+ ClassificationTrainRecord(
+ text="@foo what a cute dog!", labels=["no complaint"]
+ ),
+ ClassificationTrainRecord(
+ text="@bar this is the worst idea ever.", labels=["complaint"]
+ ),
+ ]
+ ),
+ "torch_dtype": torch.bfloat16,
+ "device": "cpu",
+ }
+ causal_lm_train_kwargs.update(patch_kwargs)
+
+ model_name = causal_lm_train_kwargs["base_model"]._model_name
+ module = caikit_nlp.modules.text_generation.PeftPromptTuning
+ with temp_config(training_data_limit={module.MODULE_ID: {model_name: 2}}):
+
+ model = module.train(**causal_lm_train_kwargs)
+ assert model
+
+
+def test_train_with_non_existent_limit_success(causal_lm_train_kwargs, set_cpu_device):
+ """Check if we are able to train successfully if training data limit doesn't exist for particular model"""
+ patch_kwargs = {
+ "num_epochs": 1,
+ "verbalizer": "Tweet text : {{input}} Label : ",
+ "train_stream": caikit.core.data_model.DataStream.from_iterable(
+ [
+ ClassificationTrainRecord(
+ text="@foo what a cute dog!", labels=["no complaint"]
+ )
+ ]
+ ),
+ "torch_dtype": torch.bfloat16,
+ "device": "cpu",
+ }
+ causal_lm_train_kwargs.update(patch_kwargs)
+
+ model_name = causal_lm_train_kwargs["base_model"]._model_name
+ module = caikit_nlp.modules.text_generation.PeftPromptTuning
+ with temp_config(training_data_limit={module.MODULE_ID: {"foo": 2}}):
+
+ model = module.train(**causal_lm_train_kwargs)
+ assert model
+
+
+def test_train_with_no_limit_for_module(causal_lm_train_kwargs, set_cpu_device):
+ """Check if we are able to train successfully if training data limit doesn't exist prompt tuning module"""
+ patch_kwargs = {
+ "num_epochs": 1,
+ "verbalizer": "Tweet text : {{input}} Label : ",
+ "train_stream": caikit.core.data_model.DataStream.from_iterable(
+ [
+ ClassificationTrainRecord(
+ text="@foo what a cute dog!", labels=["no complaint"]
+ )
+ ]
+ ),
+ "torch_dtype": torch.bfloat16,
+ "device": "cpu",
+ }
+ causal_lm_train_kwargs.update(patch_kwargs)
+
+ model_name = causal_lm_train_kwargs["base_model"]._model_name
+ module = caikit_nlp.modules.text_generation.PeftPromptTuning
+ with temp_config(training_data_limit={}):
+
+ model = module.train(**causal_lm_train_kwargs)
+ assert model