Added jsonl dump and costing

Author: ashwinprasadme

src/target_tools/llms/src/runner.py

@@ -199,6 +199,10 @@ def model_evaluation_openai(
 
     prompts = [x["prompt"] for x in id_mapping.values()]
 
+    utils.get_prompt_cost(prompts)
+    utils.dump_ft_jsonl(id_mapping, f"{results_dst}/ft_dataset.jsonl")
+    utils.dump_batch_prompt_jsonl(id_mapping, f"{results_dst}/batch_prompt.jsonl")
+
     request_outputs = openai_helpers.process_requests(
         model_name,
         prompts,
@@ -272,23 +276,31 @@ def main_runner(args, runner_config, models_to_run, openai_models_models_to_run)
             else:
                 model_path = model["lora_repo"]
 
-            pipe = transformers_helpers.load_model_and_configurations(
-                args.hf_token, model_path, model["quantization"], TEMPARATURE
-            )
-            model_start_time = time.time()
-            model_evaluation_transformers(
-                model["name"],
-                args.prompt_id,
-                python_files,
-                pipe,
-                results_dst,
-                use_system_prompt=model["use_system_prompt"],
-                batch_size=model["batch_size"],
-            )
+            pipe = None
+            try:
+                pipe = transformers_helpers.load_model_and_configurations(
+                    args.hf_token, model_path, model["quantization"], TEMPARATURE
+                )
+                model_start_time = time.time()
+                model_evaluation_transformers(
+                    model["name"],
+                    args.prompt_id,
+                    python_files,
+                    pipe,
+                    results_dst,
+                    use_system_prompt=model["use_system_prompt"],
+                    batch_size=model["batch_size"],
+                )
 
-            del pipe
-            gc.collect()
-            torch.cuda.empty_cache()
+            except Exception as e:
+                logger.error(f"Error in model {model['name']}: {e}")
+                error_count += 1
+                traceback.print_exc()
+            finally:
+                if pipe is not None:
+                    del pipe
+                gc.collect()
+                torch.cuda.empty_cache()
 
         logger.info(
             f"Model {model['name']} finished in {time.time()-model_start_time:.2f} seconds"

src/target_tools/llms/src/utils.py

@@ -9,6 +9,7 @@
 import logging
 import prompts
 import copy
+import tiktoken
 
 logger = logging.getLogger("runner")
 logger.setLevel(logging.DEBUG)
@@ -262,6 +263,58 @@ def get_prompt(prompt_id, file_path, answers_placeholders=True, use_system_promp
     return prompt
 
 
+def dump_ft_jsonl(id_mapping, output_file):
+    mappings = copy.deepcopy(id_mapping)
+    for _m in mappings.values():
+        print(_m)
+        assistant_message = {
+            "role": "assistant",
+            "content": generate_answers_for_fine_tuning(_m["json_filepath"]),
+        }
+        _m["prompt"].append(assistant_message)
+
+    prompts = [x["prompt"] for x in mappings.values()]
+
+    with open(output_file, "w") as output:
+        for _m in prompts:
+            output.write(json.dumps(_m))
+            output.write("\n")
+
+
+def dump_batch_prompt_jsonl(id_mapping, output_file):
+    prompts = [x["prompt"] for x in id_mapping.values()]
+
+    with open(output_file, "w") as output:
+        for _m in prompts:
+            output.write(json.dumps(_m))
+            output.write("\n")
+
+
+def get_prompt_cost(prompts):
+    """
+    Retrieves the token count of the given text.
+
+    Args:
+        text (str): The text to be tokenized.
+
+    Returns:
+        int: The token count.
+    """
+
+    prices_per_token = {
+        "gpt-4o": 0.000005,
+        "gpt-4o-mini": 0.00000015,
+    }
+
+    for model, price in prices_per_token.items():
+        encoding = tiktoken.encoding_for_model(model)
+        number_of_tokens = len(encoding.encode(str(prompts)))
+        logger.info(
+            f"Number of tokens for model `{model}`: {number_of_tokens}"
+            + f" Cost: {number_of_tokens * price:.5f}"
+        )
+
+
 # Example usage:
 # loader = ConfigLoader("models_config.yaml")
 # loader.load_config()