limit token

app.py

@@ -159,14 +159,16 @@ def evaluate(
     instruction,
     input=None,
     temperature=0.7,
-    max_new_tokens=256,
+    max_tokens=384,
 ):
     num_beams: int = 1
     top_p: float = 1.0
     top_k: int = 0
     prompt = generate_prompt(instruction, input)
     inputs = tokenizer(prompt, return_tensors="pt")
-    input_ids = inputs["input_ids"][:256].to(device)
+    if len(inputs["input_ids"][0]) > max_tokens:
+        return f"please reduce length. Currently, {len(inputs['input_ids'][0])} token are used.", gr.update(interactive=True), gr.update(interactive=True)
+    input_ids = inputs["input_ids"].to(device)
     generation_config = GenerationConfig(
         temperature=temperature,
         top_p=top_p,
@@ -181,7 +183,7 @@ def evaluate(
             generation_config=generation_config,
             return_dict_in_generate=True,
             output_scores=True,
-            max_new_tokens=max_new_tokens,
+            max_new_tokens=max_tokens-len(input_ids),
         )
     s = generation_output.sequences[0]
     output = tokenizer.decode(s, skip_special_tokens=True)
@@ -205,7 +207,7 @@ def evaluate(
                     "top_p": top_p,
                     "top_k": top_k,
                     "num_beams": num_beams,
-                    "max_new_tokens": max_new_tokens,
+                    "max_tokens": max_tokens,
                 },
             )
         except Exception as e:
@@ -255,13 +257,13 @@ with gr.Blocks(
                 interactive=True,
                 label="Temperature",
             )
-            max_new_tokens = gr.Slider(
-                minimum=1,
-                maximum=128,
-                value=64,
+            max_tokens = gr.Slider(
+                minimum=20,
+                maximum=384,
+                value=128,
                 step=1,
                 interactive=True,
-                label="Max length",
+                label="Max length (Pre-prompt + instruction + input + output))",
             )
 
     with gr.Column(elem_id="user_consent_container") as user_consent_block:
@@ -305,13 +307,13 @@ with gr.Blocks(
     inputs.submit(no_interactive, [], [submit_button, clear_button])
     inputs.submit(
         evaluate,
-        [instruction, inputs, temperature, max_new_tokens],
+        [instruction, inputs, temperature, max_tokens],
         [outputs, submit_button, clear_button],
     )
     submit_button.click(no_interactive, [], [submit_button, clear_button])
     submit_button.click(
         evaluate,
-        [instruction, inputs, temperature, max_new_tokens],
+        [instruction, inputs, temperature, max_tokens],
         [outputs, submit_button, clear_button],
     )
     clear_button.click(reset_textbox, [], [instruction, inputs, outputs], queue=False)