Max tokens parameter

Author: gmorgachev

src/compressa/perf/cli/__main__.py

@@ -20,7 +20,8 @@ def run_experiment_args(args):
         num_runners=args.num_runners,
         generate_prompts=args.generate_prompts,
         num_prompts=args.num_prompts,
-        prompt_length=args.prompt_length
+        prompt_length=args.prompt_length,
+        max_tokens=args.max_tokens
     )
 
 
@@ -123,6 +124,9 @@ def main():
     parser_run.add_argument(
         "--prompt_length", type=int, default=100, help="Length of each generated prompt (if --generate_prompts is used)"
     )
+    parser_run.add_argument(
+        "--max_tokens", type=int, default=1000, help="Maximum number of tokens for the model to generate"
+    )
     parser_run.set_defaults(func=run_experiment_args)
 
     parser_report = subparsers.add_parser(

src/compressa/perf/cli/tools.py

@@ -51,6 +51,8 @@ def generate_random_text(length):
         word = ''.join(random.choice(string.ascii_lowercase) for _ in range(word_length))
         words.append(word)
         current_length += len(word) + 1
+    
+    words.append(". Repeat this text at least 10 times. Number the repetitions.")
     return ' '.join(words)[:length]
 
 def generate_prompts_list(num_prompts, prompt_length):
@@ -73,7 +75,8 @@ def run_experiment(
     num_runners: int = 10,
     generate_prompts: bool = False,
     num_prompts: int = 100,
-    prompt_length: int = 100
+    prompt_length: int = 100,
+    max_tokens: int = 1000,
 ):
     if not openai_api_key:
         raise ValueError("OPENAI_API_KEY is not set")
@@ -85,7 +88,7 @@ def run_experiment(
             openai_api_key=openai_api_key,
             openai_url=openai_url,
             model_name=model_name,
-            num_runners=num_runners
+            num_runners=num_runners,
         )
 
         experiment = Experiment(
@@ -106,7 +109,8 @@ def run_experiment(
         experiment_runner.run_experiment(
             experiment_id=experiment.id,
             prompts=prompts,
-            num_tasks=num_tasks
+            num_tasks=num_tasks,
+            max_tokens=max_tokens,
         )
 
         # Run analysis after the experiment

src/compressa/perf/experiment/inference.py

@@ -38,7 +38,7 @@ def run_inference(
         self,
         experiment_id: int,
         prompt: str,
-        max_tokens: int = 1000,
+        max_tokens: int,
     ):
         start_time = time.time()
 
@@ -107,6 +107,7 @@ def store_experiment_parameters(
         self,
         experiment_id: int,
         num_tasks: int,
+        max_tokens: int,
     ):
         parameters = [
             Parameter(
@@ -127,6 +128,12 @@ def store_experiment_parameters(
                 key="openai_url",
                 value=self.openai_url,
             ),
+            Parameter(
+                id=None,
+                experiment_id=experiment_id,
+                key="max_tokens",
+                value=str(max_tokens),
+            ),
         ]
         for param in parameters:
             insert_parameter(self.conn, param)
@@ -136,6 +143,7 @@ def run_experiment(
         experiment_id: int,
         prompts: List[str],
         num_tasks: int = 100,
+        max_tokens: int = 1000,
     ):
         all_measurements = []
         with ThreadPoolExecutor(max_workers=self.num_runners) as executor:
@@ -149,7 +157,7 @@ def run_experiment(
                 for _ in range(self.num_runners)
             ]
             futures = [
-                executor.submit(runners[i % self.num_runners].run_inference, experiment_id, random.choice(prompts))
+                executor.submit(runners[i % self.num_runners].run_inference, experiment_id, random.choice(prompts), max_tokens)
                 for i in range(num_tasks)
             ]
             for future in as_completed(futures):
@@ -160,6 +168,10 @@ def run_experiment(
                 except Exception as e:
                     logger.error(f"Task failed: {e}")
 
-        self.store_experiment_parameters(experiment_id, num_tasks)
+        self.store_experiment_parameters(
+            experiment_id,
+            num_tasks,
+            max_tokens,
+        )
         for measurement in all_measurements:
             insert_measurement(self.conn, measurement)