Add long context tests and fix tokenizer truncation for activated rope_scaling

src/client_test.py

@@ -64,7 +64,7 @@ def get_client(serialize=True):
     return client
 
 
-def get_args(prompt, prompt_type, chat=False, stream_output=False,
+def get_args(prompt, prompt_type=None, chat=False, stream_output=False,
              max_new_tokens=50,
              top_k_docs=3,
              langchain_mode='Disabled',
@@ -245,9 +245,9 @@ def test_client_chat_stream(prompt_type='human_bot'):
                            langchain_agents=[])
 
 
-def run_client_chat(prompt, prompt_type, stream_output, max_new_tokens,
+def run_client_chat(prompt, stream_output, max_new_tokens,
                     langchain_mode, langchain_action, langchain_agents,
-                    prompt_dict=None):
+                    prompt_type=None, prompt_dict=None):
     client = get_client(serialize=False)
 
     kwargs, args = get_args(prompt, prompt_type, chat=True, stream_output=stream_output,

src/gen.py

@@ -1071,7 +1071,7 @@ def get_model(
         else:
             tokenizer = tokenizer_loader.from_pretrained(tokenizer_base_model, **tokenizer_kwargs)
             # sets raw (no cushion) limit
-            set_model_max_len(config, tokenizer, verbose=False)
+            set_model_max_len(config, tokenizer, verbose=False, rope_scaling=rope_scaling)
             # if using fake tokenizer, not really accurate when lots of numbers, give a bit of buffer, else get:
             # Generation Failed: Input validation error: `inputs` must have less than 2048 tokens. Given: 2233
             tokenizer.model_max_length = tokenizer.model_max_length - 50
@@ -1313,27 +1313,33 @@ def get_hf_model(load_8bit: bool = False,
         if torch.__version__ >= "2" and sys.platform != "win32" and compile_model:
             model = torch.compile(model)
 
-    set_model_max_len(config, tokenizer, verbose=False, reward_type=reward_type)
+    set_model_max_len(config, tokenizer, verbose=False, reward_type=reward_type, rope_scaling=rope_scaling)
 
     return model, tokenizer, device
 
 
-def set_model_max_len(config, tokenizer, verbose=False, reward_type=False):
+def set_model_max_len(config, tokenizer, verbose=False, reward_type=False, rope_scaling=None):
+    rope_scaling_factor = 1
+    if rope_scaling:
+        rope_scaling_factor = rope_scaling.get('factor')
+        assert isinstance(rope_scaling_factor, int)
     if reward_type:
         # limit deberta, else uses too much memory and not worth response score
         tokenizer.model_max_length = 512
     if hasattr(config, 'max_seq_len') and isinstance(config.max_seq_len, int):
-        tokenizer.model_max_length = config.max_seq_len
+        tokenizer.model_max_length = config.max_seq_len * rope_scaling_factor
     elif hasattr(config, 'max_position_embeddings') and isinstance(config.max_position_embeddings, int):
         # help automatically limit inputs to generate
-        tokenizer.model_max_length = config.max_position_embeddings
+        tokenizer.model_max_length = config.max_position_embeddings * rope_scaling_factor
     else:
         if verbose:
-            print("Could not determine model_max_length, setting to 2048", flush=True)
-        tokenizer.model_max_length = 2048
+            print(f"Could not determine model_max_length, setting to {2048 * rope_scaling_factor}", flush=True)
+        # hopefully not for Llama2 models
+        tokenizer.model_max_length = 2048 * rope_scaling_factor
     # for bug in HF transformers
     if tokenizer.model_max_length > 100000000:
-        tokenizer.model_max_length = 2048
+        # hopefully not for Llama2 models
+        tokenizer.model_max_length = 2048 * rope_scaling_factor
 
 
 def pop_unused_model_kwargs(model_kwargs):
@@ -2007,7 +2013,7 @@ def evaluate(
                                  where_from=where_from, extra_dict=extra_dict)
         return
     else:
-        assert not inference_server, "inferene_server=%s not supported" % inference_server
+        assert not inference_server, "inference_server=%s not supported" % inference_server
 
     if isinstance(tokenizer, str):
         # pipeline