test: use mock tokenizer

Author: CWHer

applications/Chat/tests/test_experience.py

@@ -11,7 +11,6 @@
 from coati.trainer.ppo import _set_default_generate_kwargs
 from coati.trainer.strategies import DDPStrategy, GeminiStrategy
 from coati.trainer.strategies.colossalai import LowLevelZeroStrategy
-from transformers import PreTrainedTokenizer
 from transformers.models.gpt2.configuration_gpt2 import GPT2Config
 
 from colossalai.testing import rerun_if_address_is_in_use, spawn
@@ -58,8 +57,13 @@ def make_and_consume_experience(strategy):
     actor, critic, initial_model, reward_model = \
         strategy.prepare(actor, critic, initial_model, reward_model)
 
-    tokenizer = PreTrainedTokenizer()
-    tokenizer.padding_side = "left"
+    class MockTokenizer():
+        def __init__(self):
+            self.padding_side = "left"
+            self.eos_token_id = 0
+            self.pad_token_id = 0
+
+    tokenizer = MockTokenizer()
     experience_maker = NaiveExperienceMaker(actor, critic, reward_model, initial_model, tokenizer)
     data_buffer = NaiveExperienceBuffer(SAMPLE_BATCH_SIZE, cpu_offload=False)
 

applications/Chat/tests/test_models.py

@@ -13,7 +13,6 @@
 from coati.models.loss import GPTLMLoss, LogExpLoss, LogSigLoss, PolicyLoss, ValueLoss
 from coati.models.opt import OPTRM, OPTActor, OPTCritic
 from coati.models.utils import calc_action_log_probs, masked_mean
-from transformers import PreTrainedTokenizer
 
 
 @pytest.mark.gpu
@@ -38,10 +37,16 @@ def test_generation(actor_maker: Callable[[], Actor],
                     seq_len: int,
                     generate_kwargs: Dict[str, Any]
                     ):
+
+    class MockTokenizer():
+        def __init__(self):
+            self.padding_side = "left"
+            self.eos_token_id = 0
+            self.pad_token_id = 0
+
     actor = actor_maker()
     input_ids = torch.randint(0, 100, (batch_size, seq_len)).cuda()
-    tokenizer = PreTrainedTokenizer()
-    tokenizer.padding_side = "left"
+    tokenizer = MockTokenizer()
     sequences = generate(actor.cuda(), input_ids, tokenizer, **generate_kwargs)
     assert sequences.shape == (batch_size, generate_kwargs["max_length"])