File size: 21,067 Bytes
42da440 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 | import hashlib
import json
import random
import traceback
from tqdm import asyncio as tqdmasyncio
import os
import sys # Add sys import
import asyncio # Add asyncio import
from augmentoolkit.generation_functions.generation_step_class import GenerationStep
from augmentoolkit.generation_functions.pipeline_step_class import (
PipelineStep,
filter_out_nonpresent_keys,
)
# TODO make this stop checking if/when it becomes impossible for either failing or passing to occur. I.e., early stopping in the event of the result being determined already.
# Also note that working with abstractions makes the logic around skipping steps NICE AND EASY
# we just put a majority vote and its "remove failed items" check into a conditional
# and that's the branching
# all nice and easy to look at
# and easier to learn from too
class MajorityVoteStep(
PipelineStep
): # Basic idea is that, result will always be a boolean, and we run this multiple times until we have a result list as long as we expect (runtime arg), and then we take the majority vote of that. Also, runtime arg of % positive required to pass to allow for easy supermajority.
def __init__(
self,
prompt_path,
regex,
sampling_params,
output_file,
output_processor,
result_key,
max_retries=3,
log_full_outputs=False,
final_determination_key=None,
vote_count_needed=1,
percent_true_to_pass=0.5,
validation_function=lambda x, y: {"result": True, "message": "default message"},
**kwargs,
):
# print(f"[DEBUG MajorityVoteStep __init__] Initializing MajorityVoteStep.")
# print(f"[DEBUG MajorityVoteStep __init__] vote_count_needed={vote_count_needed}, percent_true_to_pass={percent_true_to_pass}, result_key='{result_key}', final_determination_key='{final_determination_key}'")
self.vote_count_needed = vote_count_needed
self.percent_true_to_pass = percent_true_to_pass
self.final_determination_key = final_determination_key
super().__init__(
prompt_path=prompt_path,
regex=regex,
sampling_params=sampling_params,
output_file=output_file,
output_processor=output_processor,
result_key=result_key,
max_retries=max_retries,
log_full_outputs=log_full_outputs,
validation_function=validation_function,
**kwargs,
)
# print(f"[DEBUG MajorityVoteStep __init__] Initialization complete.")
async def read_previous_output(self, key, input_dict):
# print(f"[DEBUG MajorityVoteStep read_previous_output] Reading previous output for key='{key}'.")
entry = input_dict.get(str(key), {})
current_votes = entry.get(self.result_key, [])
# print(f"[DEBUG MajorityVoteStep read_previous_output] Found {len(current_votes)} existing votes for key='{key}'.")
if len(current_votes) >= self.vote_count_needed:
# print(f"[DEBUG MajorityVoteStep read_previous_output] Vote count needed ({self.vote_count_needed}) met for key='{key}'. Returning True.")
return True, current_votes
# print(f"[DEBUG MajorityVoteStep read_previous_output] Vote count needed ({self.vote_count_needed}) not met for key='{key}'. Returning False.")
return False, current_votes # get the votes so far
def save(
self,
result,
key,
input_dict,
input_data,
full_response,
full_input,
include_details=False,
completion_mode=False,
):
# print(f"[DEBUG MajorityVoteStep save] Saving result for key='{key}'. Result={result}, include_details={include_details}")
# Get existing entry from input_dict first, avoid copying input_data if key exists
entry = input_dict.get(str(key))
if entry is None:
# Only copy if the key is genuinely missing from input_dict
# print(f"[DEBUG MajorityVoteStep save] Key '{key}' not found in input_dict, creating new entry.")
entry = {}
if self.result_key not in entry:
# print(f"[DEBUG MajorityVoteStep save] Initializing result list for key='{key}'.")
entry[self.result_key] = []
entry[self.result_key].append(result) # result is a boolean?
# print(f"[DEBUG MajorityVoteStep save] Appended result. New vote count for key='{key}': {len(entry[self.result_key])}.")
# we hold off on saving the details possibly until the end, because we only want to save those that match the majority vote result
if (
include_details
): # well actually, we save them all but then at the end we filter... no wait what's actually the point. The logic can be shoved in the end datagen pipeline, if we keep all the boolean results together then... wait but we will not do that. That's not what happens. They all get smushed into the asme list. OK so we should filter.
# print(f"[DEBUG MajorityVoteStep save] Including details for key='{key}'.")
if not entry.get(self.details_key):
# print(f"[DEBUG MajorityVoteStep save] Initializing details list for key='{key}'.")
entry[self.details_key] = []
entry[self.details_key].append(
{
"full_response": "...", # Avoid printing large objects
"full_input": "...", # Avoid printing large objects
"completion_mode": completion_mode,
"bool_result": result,
}
)
# print(f"[DEBUG MajorityVoteStep save] Appended details. New details count for key='{key}': {len(entry[self.details_key])}.")
input_dict[str(key)] = entry
# print(f"[DEBUG MajorityVoteStep save] Saved entry for key='{key}'.")
return entry
def evaluate_final_count_and_save(
self, current_votes, input_dict, input_data, output_dir, key, include_details
):
# print(f"[DEBUG MajorityVoteStep evaluate_final_count_and_save] Evaluating final count for key='{key}'. Current vote count: {len(current_votes)}.")
true_count = sum(1 for vote in current_votes if vote)
total_votes = len(current_votes)
# print(f"[DEBUG MajorityVoteStep evaluate_final_count_and_save] True counts: {true_count}, Total votes: {total_votes}.")
if total_votes == 0:
# print(f"[DEBUG MajorityVoteStep evaluate_final_count_and_save] No votes found for key='{key}'. Returning False.")
return False
judgement = (true_count / total_votes) >= self.percent_true_to_pass
# print(f"[DEBUG MajorityVoteStep evaluate_final_count_and_save] Judgement for key='{key}': {judgement} (threshold={self.percent_true_to_pass}).")
# Entry should exist in input_dict at this point, fetch directly
entry = input_dict.get(str(key))
if entry is None:
# Fallback if key unexpectedly missing, log or handle error?
# print(f" WARNING: Key {key} unexpectedly missing from input_dict in evaluate_final_count_and_save. Creating new entry.")
entry = {} # Fallback to original logic if absolutely necessary
else:
# print(f"[DEBUG MajorityVoteStep evaluate_final_count_and_save] Found existing entry for key='{key}'.")
pass # Explicitly do nothing if entry exists
entry[self.final_determination_key] = judgement
# print(f"[DEBUG MajorityVoteStep evaluate_final_count_and_save] Set final determination '{self.final_determination_key}' for key='{key}' to {judgement}.")
# filtered detail list
if include_details:
# print(f"[DEBUG MajorityVoteStep evaluate_final_count_and_save] Filtering details for key='{key}' based on judgement ({judgement}).")
# print("INCLUDING DETAILS -- MAJORITY VOTE")
original_details_count = len(entry.get(self.details_key, []))
filtered_details = [
detail
for detail in entry.get(self.details_key, [])
if detail["bool_result"] == judgement
] # the beauty is, it does not matter if we re-run. Because then all the details will still be of the same bool result that passed the first time. So the details list is not changed.
# print([det["bool_result"] for det in filtered_details])
# print(f"[DEBUG MajorityVoteStep evaluate_final_count_and_save] Original details count: {original_details_count}, Filtered details count: {len(filtered_details)} for key='{key}'.")
entry[self.details_key] = filtered_details
input_dict[str(key)] = (
entry # where before we kept this up to date with the xisting data, now it IS the existing data. Much simpler eh?
)
# print(f"[DEBUG MajorityVoteStep evaluate_final_count_and_save] Updated entry in input_dict for key='{key}'.")
return judgement
async def run(
self,
key,
input_data,
engine_wrapper,
input_dict,
default_prompt_folder,
prompt_folder,
output_dir,
include_details,
completion_mode,
use_stop,
**kwargs,
):
# print(f"[DEBUG MajorityVoteStep run] Starting run for key='{key}'.")
full_prompt_path = (
self.prompt_path + ".yaml"
if not completion_mode
else self.prompt_path + ".txt"
)
# print(f"[DEBUG MajorityVoteStep run] Prompt path: '{full_prompt_path}'")
res, current_votes = await self.read_previous_output(key, input_dict)
# print(f"[DEBUG MajorityVoteStep run] Read previous output for key='{key}'. Result: {res}, Current votes: {len(current_votes)}.")
if res:
# print(f"[DEBUG MajorityVoteStep run] Vote count already met for key='{key}'. Evaluating final count.")
self.evaluate_final_count_and_save(
current_votes=current_votes,
input_dict=input_dict,
input_data=input_data,
output_dir=output_dir,
key=key,
include_details=include_details,
)
# print(f"[DEBUG MajorityVoteStep run] Finished run early for key='{key}' as vote count was met.")
return # Skip generation if votes are already sufficient
processed_data, additional_kwargs = self.process_input_data(input_data)
# print(f"[DEBUG MajorityVoteStep run] Processed input data for key='{key}'.")
existing_variations = []
# Check existing variations using key
output_path = self.make_output_path(output_dir)
# print(f"[DEBUG MajorityVoteStep run] Checking for existing variations in file: '{output_path}' for key='{key}'.")
try:
# print(f"[DEBUG MajorityVoteStep run] Acquired lock for '{output_path}.lock'. Reading file.")
try:
with open(output_path, "r", encoding="utf-8") as f:
existing_data = json.load(f)
existing_variations = existing_data.get(str(key), {}).get(
self.result_key, []
)
# print(f"[DEBUG MajorityVoteStep run] Loaded {len(existing_variations)} existing variations from file for key='{key}'.")
except (FileNotFoundError, json.JSONDecodeError) as e:
# print(f"[DEBUG MajorityVoteStep run] File '{output_path}' not found or JSON error for key='{key}': {e}. Starting with 0 existing variations.")
pass # No existing data is fine
except Exception as e:
print(
f"[DEBUG MajorityVoteStep run] Error during file lock/read for '{output_path}': {e}"
)
# Decide if we should continue or raise
# Generate remaining variations (it is in a loop we generate like random variations does)
needed_generations = self.vote_count_needed - len(existing_variations)
# print(f"[DEBUG MajorityVoteStep run] Need to generate {needed_generations} more votes for key='{key}' (Have {len(existing_variations)}, Need {self.vote_count_needed}).")
for i in range(len(existing_variations), self.vote_count_needed):
# print(f"[DEBUG MajorityVoteStep run] Starting generation loop iteration {i+1}/{self.vote_count_needed} for key='{key}'.")
error_message = ""
complete = False
attempt = 0
current_max_retries = self.max_retries
while not complete and attempt < current_max_retries:
attempt += 1
# print(f"[DEBUG MajorityVoteStep run] Generation attempt {attempt}/{current_max_retries} for key='{key}', iteration {i+1}.")
try:
result, full_output, full_response, full_input = (
await self.generate_data(
processed_data,
engine_wrapper,
full_prompt_path,
prompt_folder,
default_prompt_folder,
completion_mode,
use_stop,
error_message=error_message,
**kwargs,
**additional_kwargs,
)
)
# print(f"[DEBUG MajorityVoteStep run] generate_data returned for key='{key}'. Result: {result}.")
validation_result = self.validation_function(result, input_data)
if validation_result["result"]:
# print(f"[DEBUG MajorityVoteStep run] Validation function passed for key='{key}'.")
complete = True
else:
# print(f"[DEBUG MajorityVoteStep run] Validation function failed for key='{key}'.")
error_message = validation_result["message"]
pass # Explicitly do nothing if validation fails
except Exception as e:
# print(f"[DEBUG MajorityVoteStep run] Exception during generate_data for key='{key}', attempt {attempt}: {e}")
error_message = str(e)
traceback.print_exc()
if not complete:
# print(f"[DEBUG MajorityVoteStep run] Failed to generate a valid result after {current_max_retries} attempts for key='{key}'. Returning.")
return # Stop processing this key if generation failed
# print(f"[DEBUG MajorityVoteStep run] Generation successful for key='{key}', iteration {i+1}. Saving result.")
self.save(
result=result,
key=key,
input_dict=input_dict,
input_data=input_data, # Passing input_data, but save doesn't seem to use it directly
full_response=full_response, # Passing potentially large object
full_input=full_input, # Passing potentially large object
include_details=include_details,
completion_mode=completion_mode,
)
# print(f"[DEBUG MajorityVoteStep run] Saved result for key='{key}', iteration {i+1}.")
# print(f"[DEBUG MajorityVoteStep run] Finished generation loop for key='{key}'.")
# Re-check current votes after potential generation
res, current_votes = await self.read_previous_output(key, input_dict)
# print(f"[DEBUG MajorityVoteStep run] Re-read previous output for key='{key}' after generation loop. Result: {res}, Current votes: {len(current_votes)}.")
if res: # Should be true now if loop completed successfully
# print(f"[DEBUG MajorityVoteStep run] Evaluating final count after generation for key='{key}'.")
self.evaluate_final_count_and_save(
current_votes=current_votes,
input_dict=input_dict,
input_data=input_data,
output_dir=output_dir,
key=key,
include_details=include_details,
)
else:
# print(f"[DEBUG MajorityVoteStep run] WARNING: Vote count still not met after generation loop for key='{key}'. This shouldn't happen if loop completed.")
pass # Explicitly do nothing if vote count not met
# print(f"[DEBUG MajorityVoteStep run] Finished run for key='{key}'.")
# This way it makes it even easierto code pipelinesteps too! Just need to manipulte the dict, no files.
async def execute_pipeline(
self,
engine_wrapper,
rtwl,
default_prompt_folder,
prompt_folder,
output_dir,
completion_mode,
use_stop,
include_details,
input_dict={},
**kwargs,
):
# print(f"[DEBUG MajorityVoteStep execute_pipeline] Starting execute_pipeline. Output dir: '{output_dir}'. Initial input_dict size: {len(input_dict)}.")
# print(f"[DEBUG MajorityVoteStep execute_pipeline] Loading dataset...")
self.load_dataset(input_dict=input_dict, output_dir=output_dir)
# print(f"[DEBUG MajorityVoteStep execute_pipeline] Dataset loaded. input_dict size: {len(input_dict)}.")
num_tasks = len(input_dict)
# print(f"[DEBUG MajorityVoteStep execute_pipeline] Creating {num_tasks} run tasks.")
try:
data_generations_tasks = [
self.run(
key=key,
input_data=value,
engine_wrapper=engine_wrapper,
input_dict=input_dict,
default_prompt_folder=default_prompt_folder,
prompt_folder=prompt_folder,
output_dir=output_dir,
completion_mode=completion_mode,
use_stop=use_stop,
include_details=include_details,
**kwargs,
)
for key, value in input_dict.items()
]
coroutines = [rtwl(task) for task in data_generations_tasks]
# print(f"[DEBUG MajorityVoteStep execute_pipeline] Submitting {len(coroutines)} tasks to rate limited wrapper.")
# Adding progress tracking
TASK_TIMEOUT_SECONDS = 600 # 10 minutes timeout
completed_tasks = 0
num_tasks = len(coroutines) # Use pre-calculated num_tasks
for future in tqdmasyncio.tqdm.as_completed(coroutines):
try:
# Wrap await future with timeout
await asyncio.wait_for(future, timeout=TASK_TIMEOUT_SECONDS)
completed_tasks += 1
# if completed_tasks % 100 == 0 or completed_tasks == num_tasks: # Print every 100 tasks or on the last task
# print(f"[DEBUG MajorityVoteStep execute_pipeline] Completed {completed_tasks}/{num_tasks} tasks.")
except asyncio.TimeoutError:
# completed_tasks might not be perfectly accurate here if timeout happens before increment
print(
f"\nWARNING: Task (approx {completed_tasks+1}/{num_tasks}) for step '{self.output_file}' timed out after {TASK_TIMEOUT_SECONDS} seconds.",
file=sys.stderr,
)
# Task is cancelled by wait_for, loop continues
except Exception as task_exception:
# print(f"[DEBUG MajorityVoteStep execute_pipeline] Exception occurred in awaited task: {task_exception}")
traceback.print_exc() # Print traceback for individual task failure
# print(f"[DEBUG MajorityVoteStep execute_pipeline] All tasks completed.")
except Exception as e:
# print(f"[DEBUG MajorityVoteStep execute_pipeline] Exception occurred during task execution loop: {e}")
traceback.print_exc()
# raise e # Decide if outer exception should halt everything
finally:
# print(f"[DEBUG MajorityVoteStep execute_pipeline] Entering finally block.")
# print(f"[DEBUG MajorityVoteStep execute_pipeline] Saving dataset. Current input_dict size: {len(input_dict)}.")
self.save_dataset(input_dict=input_dict, output_dir=output_dir)
# print(f"[DEBUG MajorityVoteStep execute_pipeline] Dataset saved.")
# print(f"[DEBUG MajorityVoteStep execute_pipeline] Filtering out non-present keys using key '{self.result_key}'. Current input_dict size: {len(input_dict)}.")
filter_out_nonpresent_keys(input_dict, key_to_check=self.result_key)
# print(f"[DEBUG MajorityVoteStep execute_pipeline] Filtering complete. Final input_dict size: {len(input_dict)}.")
# print(f"[DEBUG MajorityVoteStep execute_pipeline] Finished execute_pipeline.")
|