accessible_worlds.aw_generate_concurrent
1import asyncio 2from pathlib import Path 3import time 4import argparse 5from typing import Any 6from dataclasses import asdict 7 8import numpy as np 9 10from openai import AsyncOpenAI 11 12from .aw_vocabulary import AwVocabulary 13 14from .aw_build_prompt import ( 15 assemble_system_prompt, 16 add_core_integrations, 17 add_moral_integrations, 18 add_underrepresented_integrations, 19 add_recursive_integrations 20) 21 22from .aw_requests import get_response 23from .aw_config import AwConfig 24 25from .aw_utils import ( 26 safe_save, 27 load_json, 28 serialize_np_rng_state, 29 deserialize_np_rng_state, 30 BLUE, 31 GREEN, 32 YELLOW, 33 ORANGE, 34 PINK, 35 RESET, 36 split_and_preserve_names, 37 preprocess_passage 38) 39 40from .aw_sentence_refinement import refine_sentences 41 42async def generate_batch( 43 semaphore, 44 client, 45 model : str, 46 system_prompt : str, 47 refinement_system_prompt : str, 48 batch : list[dict[str,Any]], 49 vocabulary : AwVocabulary, 50 known_names : set[str], 51 refine : bool, 52 refine_at : int, 53 retries_on_new : int, 54 retry_on_new_at: int, 55 retry_on_new_form_at : int, 56 rand_seed : int, 57 max_length : int, 58 min_length : int, 59 max_request_attempts : int = 3 60) -> list[ tuple[str, list[str], list[tuple[str,str]], int, int ]] : 61 62 async def generate_one( idx : int ) -> tuple[str, list[str], list[tuple[str,str]], int, int] : 63 64 new_words = set() 65 completion = "" 66 words = [] 67 refinements = [] 68 69 # If we have reached retry_on_new_at, then we will regenerate up to max( max_attempts, 1 ) times until 70 # we get a generation that contains no 'new' words 71 # Note that a 'non-new' word may be a new form of an existing word, which we allow 72 # So the vocabulary may still grow even if a 'new' word is not found 73 74 def get_new_words( words : list[str] ) -> tuple[set[str], set[str]] : 75 76 completion_word_frequencies = vocabulary.get_word_counts( words ) 77 # we consider a word 'new' if no form of the word is a core word or already in the vocabulary 78 79 new_words = set( [ 80 word for word, count in completion_word_frequencies.items() 81 if count == 0 and not vocabulary.is_form_of_core_word(word) and not vocabulary.is_form_of_vocab_word(word) 82 ] ) 83 84 new_forms = set( [ 85 word for word, count in completion_word_frequencies.items() 86 if count == 0 87 ] ) - new_words 88 89 return new_words, new_forms 90 91 instance_idx = batch[idx]['index'] 92 att_idx = 0 93 while True : 94 95 refinements = [] 96 97 completion = await get_response( 98 semaphore=semaphore, 99 client=client, 100 model=model, 101 system_prompt=system_prompt, 102 user_prompt=batch[idx]['prompt'], 103 seed=( att_idx + rand_seed + instance_idx * 0x9e3779b9 ) % (2**32), 104 temperature=1.0, 105 reasoning_budget=0, 106 max_request_attempts=max_request_attempts 107 ) 108 109 passage = preprocess_passage( completion ) 110 111 # Names are capitalized and all other words lowercase 112 words, is_valid = split_and_preserve_names( passage=passage, vocabulary=vocabulary, known_names=known_names ) 113 new_words, new_forms = get_new_words( words=words ) 114 115 # If the final completion has 'new' words and refinement is engaged, ask the model 116 # to potentially reuse and existing one, or reconsider the existing wording 117 118 if is_valid and new_words and refine and instance_idx >= refine_at : 119 120 completion, refinements = await refine_sentences( 121 semaphore=semaphore, 122 client=client, 123 model=model, 124 system_prompt=refinement_system_prompt, 125 completion=completion, 126 new_words=new_words, 127 vocabulary=vocabulary, 128 current_idx=instance_idx, 129 rand_seed=rand_seed 130 ) 131 132 passage = preprocess_passage( completion ) 133 words, is_valid = split_and_preserve_names( passage=passage, vocabulary=vocabulary, known_names=known_names ) 134 new_words, new_forms = get_new_words( words=words ) 135 136 att_idx += 1 137 138 if len( words ) > max_length : 139 is_valid = False 140 print( f"Passage too long {len(words)}" ) 141 142 elif len( words ) < min_length : 143 is_valid = False 144 print( f"Passage too short {len(words)}" ) 145 146 retry_because_new_word = new_words and instance_idx >= retry_on_new_at and att_idx <= retries_on_new + 1 147 retry_because_new_form = new_forms and instance_idx >= retry_on_new_form_at and att_idx <= retries_on_new + 1 148 149 retry = retry_because_new_word or retry_because_new_form or not is_valid 150 151 if not retry : 152 break 153 154 if not is_valid : 155 print( f"Invalid" ) 156 157 if retry_because_new_word : 158 print( f"New words past retry on new words at: {new_words}" ) 159 160 if retry_because_new_form : 161 print( f"New word forms past retry on new words at: {new_forms}" ) 162 163 print( "Retrying ..." ) 164 print( "---------------------------------------" ) 165 166 if new_words : 167 print( f"Recording new words {new_words}" ) 168 169 if new_forms : 170 print( f"Recording new forms of existing words {new_forms}" ) 171 172 vocabulary.update_words( words ) 173 174 return passage, words, refinements, vocabulary.size(), att_idx 175 176 async with asyncio.TaskGroup() as tg: 177 tasks = [ 178 tg.create_task(generate_one(idx)) 179 for idx, _ in enumerate(batch) 180 ] 181 182 instances = [task.result() for task in tasks] 183 184 return instances 185 186async def generate( 187 api_url : str, 188 api_key : str, 189 model : str, 190 config : AwConfig, 191 output_dir : Path, 192 save_every : int, 193 concurrency : int, 194 batch_size : int, 195 offset : int, 196 n_complete : int, 197 file_idx : int, 198 world : dict[str,Any], 199 exposition_setting : dict[str,Any], 200 verbose : bool = False ) : 201 202 semaphore = asyncio.Semaphore( concurrency ) 203 204 # Initialize and Prepare state 205 206 client = AsyncOpenAI( 207 base_url=api_url, 208 api_key=api_key 209 ) 210 211 rand_seed = config.rand_seed 212 np_rng = np.random.default_rng( rand_seed ) 213 214 output_path = output_dir / f"data_{file_idx}.json" 215 vocabulary_path = output_dir / f"vocabulary.json" 216 random_state_path = output_dir / f"random_state.json" 217 218 # data file init or resume ----------------------------------- 219 220 if output_path.exists() : 221 print( f"\nResuming from {output_path}... \n" ) 222 data = load_json( output_path ) 223 instances = data[ 'instances' ] 224 time_spent = data[ 'time_spent' ] 225 total_words = data[ 'total_words' ] 226 227 # Early exit, file is already complete 228 if len( instances ) >= n_complete : 229 return 230 231 # happens if we just begun starting a new file 232 # need to get total time spent and total words from the previous file 233 elif file_idx > 0 : 234 previous_file_path = output_dir / f"data_{file_idx-1}.json" 235 if not previous_file_path.exists() : 236 raise Exception( "Missing file {previous_file_path}" ) 237 previous_data = load_json( previous_file_path ) 238 time_spent = previous_data[ 'time_spent' ] 239 total_words = previous_data[ 'total_words' ] 240 instances = [] 241 242 # happens on the first run on the first file 243 else : 244 instances = [] 245 time_spent = 0 246 total_words = 0 247 248 # Vocabulary init or resume----------------- 249 250 if vocabulary_path.exists() : 251 vocabulary = AwVocabulary.from_file( vocabulary_path ) 252 else : 253 vocabulary = AwVocabulary( 254 core_things_and_concepts=world[ 'core things and concepts' ], 255 moral_concepts=world[ 'moral concepts' ] 256 ) 257 258 # Random state init or resume----------------- 259 260 if random_state_path.exists() : 261 rs_data = load_json( random_state_path ) 262 np_rng_state = deserialize_np_rng_state( rs_data[ 'np_rng_state' ] ) 263 np_rng = np.random.default_rng() 264 np_rng.bit_generator.state = np_rng_state 265 else : 266 np_rng = np.random.default_rng( rand_seed ) 267 np_rng_state = np_rng.bit_generator.state 268 269 # ---------------------------------------------------------- 270 271 known_names = set( exposition_setting[ 'name suggestions' ] ) 272 273 # We only care about singular versions 274 name_counts = vocabulary.get_word_counts( [ name for name in known_names ] ) 275 276 system_prompt = assemble_system_prompt( 277 core_prompt=exposition_setting[ 'core system prompt' ], 278 core_elements=world[ 'core things and concepts' ], 279 exposition_setting=exposition_setting, 280 np_rng=np_rng, 281 balance_names=True, 282 name_counts=name_counts 283 ) 284 285 # Add the plural forms 286 known_names = known_names | { kn + 's' for kn in known_names } 287 288 if verbose : 289 print( f"\n{BLUE}SYSTEM PROMPT:{RESET} {system_prompt}\n" ) 290 291 # ------------------------------------------------------------ 292 293 start_idx = len( instances ) + offset 294 295 batch_idx = 0 296 batch = [] 297 end_idx = min( offset + n_complete, config.total_completions ) 298 batch_words = 0 299 300 for current_idx in range( start_idx, end_idx ) : 301 302 # core user prompt with style and concept hints 303 304 explanitory_style = np_rng.choice( exposition_setting[ 'exposition styles' ], size=1 )[0] 305 concept_choice = np_rng.choice( exposition_setting[ 'concept choices' ], size=1 )[0] 306 307 user_prompt = f"{explanitory_style} {concept_choice}" 308 309 # enrich with integrations 310 311 core_integrations : list[str] = [] 312 moral_integrations : list[str] = [] 313 recursive_integration : str = "" 314 underrepresented_integrations : list[str] = [] 315 316 if np_rng.random() < config.core_integration_prob : 317 318 user_prompt, core_integrations = add_core_integrations( 319 user_prompt, 320 vocabulary=vocabulary, 321 min_k=config.min_core_integrations, 322 max_k=config.max_core_integrations, 323 min_extra_creature_integrations=config.min_extra_creature_integrations, 324 max_extra_creature_integrations=config.max_extra_creature_integrations, 325 np_rng=np_rng 326 ) 327 328 vocabulary.update_core_integrations( core_integrations ) 329 330 if ( np_rng.random() < config.underrepresented_integration_prob and 331 current_idx >= config.begin_underrepresented_integration_at 332 ): 333 user_prompt, underrepresented_integrations = add_underrepresented_integrations( 334 user_prompt, 335 vocabulary=vocabulary, 336 min_k=config.min_underrepresented_integrations, 337 max_k=config.max_underrepresented_integrations, 338 np_rng=np_rng 339 ) 340 341 # Note, on new file instances starts empty again 342 if ( np_rng.random() < config.recursive_integration_prob and 343 current_idx >= config.begin_recursive_integration_at and 344 len( instances ) > 20 345 ): 346 347 user_prompt, recursive_integration = add_recursive_integrations( 348 user_prompt, 349 previous_passages=[ instance[ 'completion' ] for instance in instances ], 350 np_rng=np_rng 351 ) 352 353 if np_rng.random() < config.moral_integration_prob : 354 user_prompt, moral_integrations = add_moral_integrations( 355 user_prompt, 356 vocabulary=vocabulary, 357 min_k=config.min_moral_integrations, 358 max_k=config.max_moral_integrations, 359 np_rng=np_rng 360 361 ) 362 vocabulary.update_moral_integrations( moral_integrations ) 363 364 user_prompt += f" {exposition_setting[ 'user prompt reminder' ]}" 365 366 batch.append( { 367 "explanitory_style": explanitory_style, 368 "concept_choice": concept_choice, 369 "core_integrations" : core_integrations, 370 "moral_integrations" : moral_integrations, 371 "recursive_integration" : recursive_integration, 372 "underrepresented_integrations" : underrepresented_integrations, 373 "prompt": user_prompt, 374 "index" : current_idx, 375 "refinements" : [] 376 } ) 377 378 if len( batch ) == batch_size or current_idx == end_idx - 1 : 379 380 start = time.time() 381 vsz_before = vocabulary.size() 382 n_tried = 0 383 384 batch_instances = await generate_batch( 385 semaphore=semaphore, 386 client=client, 387 model=model, 388 system_prompt=system_prompt, 389 refinement_system_prompt=exposition_setting[ 'refinement system prompt' ], 390 batch=batch, 391 vocabulary=vocabulary, 392 known_names=known_names, 393 refine=config.refine_vocabulary, 394 refine_at=config.begin_vocabulary_refinement_at, 395 retries_on_new=config.retries_on_new_word, 396 retry_on_new_at=config.retry_on_new_words_at, 397 retry_on_new_form_at=config.retry_on_new_word_forms_at, 398 rand_seed=rand_seed, 399 max_length=config.max_length, 400 min_length=config.min_length, 401 max_request_attempts=10 402 ) 403 404 for idx, c in enumerate( batch_instances ) : 405 406 instance = batch[ idx ] 407 408 text = c[0] 409 words = c[1] 410 refinements = c[2] 411 new_vocab_size = c[3] 412 n_tried += c[4] 413 414 instance[ 'completion' ] = text 415 instance[ 'vocab_size' ] = new_vocab_size 416 instance[ 'refinements'] = refinements 417 418 total_words += len( words ) 419 batch_words += len( words ) 420 421 # Update the integration trackers 422 423 instances.append( instance ) 424 425 if verbose : 426 print( f"-----------------------------------------\n\n{ORANGE}INSTANCE:{RESET} {instance['index']}\n" ) 427 print( f"{GREEN}PROMPT:{RESET} {instance['prompt']}\n" ) 428 for s, r in refinements : 429 print( f"{PINK}SENTENCE REFINEMENT:{RESET}\n{s}->\n{r}\n" ) 430 print( f"{YELLOW}COMPLETION:{RESET} {text}\n" ) 431 432 np_rng_state = np_rng.bit_generator.state 433 n_completed = current_idx+1 434 tc = time.time() - start 435 time_spent += tc 436 avg_time = time_spent / n_completed 437 avg_words = total_words / n_completed 438 avg_words_per_s = total_words / time_spent 439 vsz_now = vocabulary.size() 440 n_new_words = vsz_now - vsz_before 441 442 exp_boosts = ( 0 if current_idx < config.begin_underrepresented_integration_at 443 else len(batch) * config.underrepresented_integration_prob * ( config.max_underrepresented_integrations + config.min_underrepresented_integrations ) / 2.0 444 ) 445 446 print( f"retries: {n_tried-batch_size}" ) 447 print( f"batch w/s: {batch_words/tc:.2f}" ) 448 print( f"avg w/s: {avg_words_per_s:.2f}" ) 449 print( f"batch s/cmpl: {tc/batch_size:.3f}" ) 450 print( f"avg s/cmpl: {avg_time:.3f}" ) 451 print( f"avg w/cmpl: {avg_words:.2f}" ) 452 print( f"total c: {n_completed}" ) 453 print( f"total w: {total_words}, ~{round(total_words/0.75)}tk" ) 454 print( f"vocab size: {vsz_now}" ) 455 print( f"new words: {n_new_words}" ) 456 print( f"exp boosts: {exp_boosts}\n" ) 457 458 if ( batch_idx + 1 ) % save_every == 0 : 459 460 safe_save( { 461 'instances' : instances, 462 "config" : config, 463 "np_rng_state" : serialize_np_rng_state( dict(np_rng_state) ), 464 "random_seed" : rand_seed, 465 "time_spent" : time_spent, 466 'total_words' : total_words 467 }, filepath=output_path ) 468 469 safe_save( 470 { "np_rng_state" : serialize_np_rng_state( dict(np_rng_state) ) }, 471 filepath=random_state_path 472 ) 473 474 vocabulary.to_file( vocabulary_path ) 475 476 print( "Saved.\n" ) 477 478 batch = [] 479 batch_idx += 1 480 batch_words = 0 481 482 safe_save( { 483 'instances' : instances, 484 "config" : config, 485 "np_rng_state" : serialize_np_rng_state( dict(np_rng_state) ), 486 "random_seed" : rand_seed, 487 "time_spent" : time_spent, 488 'total_words' : total_words 489 }, filepath=output_path ) 490 491 safe_save( 492 { "np_rng_state" : serialize_np_rng_state( dict(np_rng_state) ) }, 493 filepath=random_state_path 494 ) 495 496 vocabulary.to_file( vocabulary_path ) 497 498if __name__ == "__main__": 499 500 parser = argparse.ArgumentParser( description="Generate Simpler World training data." ) 501 502 parser.add_argument( "--api_url", type=str ) 503 parser.add_argument( "--api_key", type=str ) 504 parser.add_argument( "--model", type=str ) 505 parser.add_argument( "--config_path", type=str, help="Path to the SW JSON configuration file" ) 506 parser.add_argument( "--output_dir", type=str, help="Output directory" ) 507 parser.add_argument( "--save_every", type=int, default=500, help="Save frequency (default: 500)" ) 508 parser.add_argument( "--concurrency", type=int, default=1, help="Number of concurrent requests (default: 1)" ) 509 parser.add_argument( "--batch_size", type=int, default=1, help="Batch size (default: 1)" ) 510 parser.add_argument( "--offset", type=int, default=0, help="Global instance offset" ) 511 parser.add_argument( "--n_complete", type=int, help="Number of completions to generate for this pass" ) 512 parser.add_argument( "--file_idx", type=int, help="File number." ) 513 parser.add_argument( "--verbose", action='store_true', help='Enable verbose output') 514 515 args = parser.parse_args() 516 517 api_url = args.api_url 518 api_key = args.api_key 519 model = args.model 520 config_path = args.config_path 521 output_dir = args.output_dir 522 save_every = args.save_every 523 concurrency = args.concurrency 524 batch_size = args.batch_size 525 offset = args.offset 526 n_complete = args.n_complete 527 file_idx = args.file_idx 528 verbose = args.verbose 529 530 config = AwConfig( **load_json( config_path ) ) 531 532 config_output_path = Path( output_dir ) / "config.json" 533 if not config_output_path.exists() : 534 safe_save( data=asdict(config), filepath=Path( output_dir ) / "config.json" ) 535 536 this_dir = Path(__file__).parent 537 538 # Load from default if local world config don't exist yet 539 540 output_world_path = Path( output_dir ) / "world.json" 541 if not output_world_path.exists() : 542 world = load_json( this_dir / 'worlds' / ( config.world + ".json" ) ) 543 safe_save( world, filepath=output_world_path ) 544 else : 545 world = load_json( output_world_path ) 546 547 # Load from default if local exposition settings don't exist yet 548 549 exposition_output_path = Path( output_dir ) / "exposition_settings.json" 550 if not exposition_output_path.exists() : 551 exposition_setting = load_json( this_dir / 'exposition_settings' / ( config.exposition_setting + ".json" ) ) 552 safe_save( exposition_setting, filepath=exposition_output_path ) 553 else : 554 exposition_setting = load_json( exposition_output_path ) 555 556 asyncio.run( generate( 557 api_url=api_url, 558 api_key=api_key, 559 model=model, 560 config=config, 561 output_dir=Path(output_dir), 562 save_every=save_every, 563 concurrency=concurrency, 564 batch_size=batch_size, 565 offset=offset, 566 n_complete=n_complete, 567 file_idx=file_idx, 568 world=world, 569 exposition_setting=exposition_setting, 570 verbose=verbose 571 ) )
async def
generate_batch( semaphore, client, model: str, system_prompt: str, refinement_system_prompt: str, batch: list[dict[str, typing.Any]], vocabulary: accessible_worlds.aw_vocabulary.AwVocabulary, known_names: set[str], refine: bool, refine_at: int, retries_on_new: int, retry_on_new_at: int, retry_on_new_form_at: int, rand_seed: int, max_length: int, min_length: int, max_request_attempts: int = 3) -> list[tuple[str, list[str], list[tuple[str, str]], int, int]]:
43async def generate_batch( 44 semaphore, 45 client, 46 model : str, 47 system_prompt : str, 48 refinement_system_prompt : str, 49 batch : list[dict[str,Any]], 50 vocabulary : AwVocabulary, 51 known_names : set[str], 52 refine : bool, 53 refine_at : int, 54 retries_on_new : int, 55 retry_on_new_at: int, 56 retry_on_new_form_at : int, 57 rand_seed : int, 58 max_length : int, 59 min_length : int, 60 max_request_attempts : int = 3 61) -> list[ tuple[str, list[str], list[tuple[str,str]], int, int ]] : 62 63 async def generate_one( idx : int ) -> tuple[str, list[str], list[tuple[str,str]], int, int] : 64 65 new_words = set() 66 completion = "" 67 words = [] 68 refinements = [] 69 70 # If we have reached retry_on_new_at, then we will regenerate up to max( max_attempts, 1 ) times until 71 # we get a generation that contains no 'new' words 72 # Note that a 'non-new' word may be a new form of an existing word, which we allow 73 # So the vocabulary may still grow even if a 'new' word is not found 74 75 def get_new_words( words : list[str] ) -> tuple[set[str], set[str]] : 76 77 completion_word_frequencies = vocabulary.get_word_counts( words ) 78 # we consider a word 'new' if no form of the word is a core word or already in the vocabulary 79 80 new_words = set( [ 81 word for word, count in completion_word_frequencies.items() 82 if count == 0 and not vocabulary.is_form_of_core_word(word) and not vocabulary.is_form_of_vocab_word(word) 83 ] ) 84 85 new_forms = set( [ 86 word for word, count in completion_word_frequencies.items() 87 if count == 0 88 ] ) - new_words 89 90 return new_words, new_forms 91 92 instance_idx = batch[idx]['index'] 93 att_idx = 0 94 while True : 95 96 refinements = [] 97 98 completion = await get_response( 99 semaphore=semaphore, 100 client=client, 101 model=model, 102 system_prompt=system_prompt, 103 user_prompt=batch[idx]['prompt'], 104 seed=( att_idx + rand_seed + instance_idx * 0x9e3779b9 ) % (2**32), 105 temperature=1.0, 106 reasoning_budget=0, 107 max_request_attempts=max_request_attempts 108 ) 109 110 passage = preprocess_passage( completion ) 111 112 # Names are capitalized and all other words lowercase 113 words, is_valid = split_and_preserve_names( passage=passage, vocabulary=vocabulary, known_names=known_names ) 114 new_words, new_forms = get_new_words( words=words ) 115 116 # If the final completion has 'new' words and refinement is engaged, ask the model 117 # to potentially reuse and existing one, or reconsider the existing wording 118 119 if is_valid and new_words and refine and instance_idx >= refine_at : 120 121 completion, refinements = await refine_sentences( 122 semaphore=semaphore, 123 client=client, 124 model=model, 125 system_prompt=refinement_system_prompt, 126 completion=completion, 127 new_words=new_words, 128 vocabulary=vocabulary, 129 current_idx=instance_idx, 130 rand_seed=rand_seed 131 ) 132 133 passage = preprocess_passage( completion ) 134 words, is_valid = split_and_preserve_names( passage=passage, vocabulary=vocabulary, known_names=known_names ) 135 new_words, new_forms = get_new_words( words=words ) 136 137 att_idx += 1 138 139 if len( words ) > max_length : 140 is_valid = False 141 print( f"Passage too long {len(words)}" ) 142 143 elif len( words ) < min_length : 144 is_valid = False 145 print( f"Passage too short {len(words)}" ) 146 147 retry_because_new_word = new_words and instance_idx >= retry_on_new_at and att_idx <= retries_on_new + 1 148 retry_because_new_form = new_forms and instance_idx >= retry_on_new_form_at and att_idx <= retries_on_new + 1 149 150 retry = retry_because_new_word or retry_because_new_form or not is_valid 151 152 if not retry : 153 break 154 155 if not is_valid : 156 print( f"Invalid" ) 157 158 if retry_because_new_word : 159 print( f"New words past retry on new words at: {new_words}" ) 160 161 if retry_because_new_form : 162 print( f"New word forms past retry on new words at: {new_forms}" ) 163 164 print( "Retrying ..." ) 165 print( "---------------------------------------" ) 166 167 if new_words : 168 print( f"Recording new words {new_words}" ) 169 170 if new_forms : 171 print( f"Recording new forms of existing words {new_forms}" ) 172 173 vocabulary.update_words( words ) 174 175 return passage, words, refinements, vocabulary.size(), att_idx 176 177 async with asyncio.TaskGroup() as tg: 178 tasks = [ 179 tg.create_task(generate_one(idx)) 180 for idx, _ in enumerate(batch) 181 ] 182 183 instances = [task.result() for task in tasks] 184 185 return instances
async def
generate( api_url: str, api_key: str, model: str, config: accessible_worlds.aw_config.AwConfig, output_dir: pathlib.Path, save_every: int, concurrency: int, batch_size: int, offset: int, n_complete: int, file_idx: int, world: dict[str, typing.Any], exposition_setting: dict[str, typing.Any], verbose: bool = False):
187async def generate( 188 api_url : str, 189 api_key : str, 190 model : str, 191 config : AwConfig, 192 output_dir : Path, 193 save_every : int, 194 concurrency : int, 195 batch_size : int, 196 offset : int, 197 n_complete : int, 198 file_idx : int, 199 world : dict[str,Any], 200 exposition_setting : dict[str,Any], 201 verbose : bool = False ) : 202 203 semaphore = asyncio.Semaphore( concurrency ) 204 205 # Initialize and Prepare state 206 207 client = AsyncOpenAI( 208 base_url=api_url, 209 api_key=api_key 210 ) 211 212 rand_seed = config.rand_seed 213 np_rng = np.random.default_rng( rand_seed ) 214 215 output_path = output_dir / f"data_{file_idx}.json" 216 vocabulary_path = output_dir / f"vocabulary.json" 217 random_state_path = output_dir / f"random_state.json" 218 219 # data file init or resume ----------------------------------- 220 221 if output_path.exists() : 222 print( f"\nResuming from {output_path}... \n" ) 223 data = load_json( output_path ) 224 instances = data[ 'instances' ] 225 time_spent = data[ 'time_spent' ] 226 total_words = data[ 'total_words' ] 227 228 # Early exit, file is already complete 229 if len( instances ) >= n_complete : 230 return 231 232 # happens if we just begun starting a new file 233 # need to get total time spent and total words from the previous file 234 elif file_idx > 0 : 235 previous_file_path = output_dir / f"data_{file_idx-1}.json" 236 if not previous_file_path.exists() : 237 raise Exception( "Missing file {previous_file_path}" ) 238 previous_data = load_json( previous_file_path ) 239 time_spent = previous_data[ 'time_spent' ] 240 total_words = previous_data[ 'total_words' ] 241 instances = [] 242 243 # happens on the first run on the first file 244 else : 245 instances = [] 246 time_spent = 0 247 total_words = 0 248 249 # Vocabulary init or resume----------------- 250 251 if vocabulary_path.exists() : 252 vocabulary = AwVocabulary.from_file( vocabulary_path ) 253 else : 254 vocabulary = AwVocabulary( 255 core_things_and_concepts=world[ 'core things and concepts' ], 256 moral_concepts=world[ 'moral concepts' ] 257 ) 258 259 # Random state init or resume----------------- 260 261 if random_state_path.exists() : 262 rs_data = load_json( random_state_path ) 263 np_rng_state = deserialize_np_rng_state( rs_data[ 'np_rng_state' ] ) 264 np_rng = np.random.default_rng() 265 np_rng.bit_generator.state = np_rng_state 266 else : 267 np_rng = np.random.default_rng( rand_seed ) 268 np_rng_state = np_rng.bit_generator.state 269 270 # ---------------------------------------------------------- 271 272 known_names = set( exposition_setting[ 'name suggestions' ] ) 273 274 # We only care about singular versions 275 name_counts = vocabulary.get_word_counts( [ name for name in known_names ] ) 276 277 system_prompt = assemble_system_prompt( 278 core_prompt=exposition_setting[ 'core system prompt' ], 279 core_elements=world[ 'core things and concepts' ], 280 exposition_setting=exposition_setting, 281 np_rng=np_rng, 282 balance_names=True, 283 name_counts=name_counts 284 ) 285 286 # Add the plural forms 287 known_names = known_names | { kn + 's' for kn in known_names } 288 289 if verbose : 290 print( f"\n{BLUE}SYSTEM PROMPT:{RESET} {system_prompt}\n" ) 291 292 # ------------------------------------------------------------ 293 294 start_idx = len( instances ) + offset 295 296 batch_idx = 0 297 batch = [] 298 end_idx = min( offset + n_complete, config.total_completions ) 299 batch_words = 0 300 301 for current_idx in range( start_idx, end_idx ) : 302 303 # core user prompt with style and concept hints 304 305 explanitory_style = np_rng.choice( exposition_setting[ 'exposition styles' ], size=1 )[0] 306 concept_choice = np_rng.choice( exposition_setting[ 'concept choices' ], size=1 )[0] 307 308 user_prompt = f"{explanitory_style} {concept_choice}" 309 310 # enrich with integrations 311 312 core_integrations : list[str] = [] 313 moral_integrations : list[str] = [] 314 recursive_integration : str = "" 315 underrepresented_integrations : list[str] = [] 316 317 if np_rng.random() < config.core_integration_prob : 318 319 user_prompt, core_integrations = add_core_integrations( 320 user_prompt, 321 vocabulary=vocabulary, 322 min_k=config.min_core_integrations, 323 max_k=config.max_core_integrations, 324 min_extra_creature_integrations=config.min_extra_creature_integrations, 325 max_extra_creature_integrations=config.max_extra_creature_integrations, 326 np_rng=np_rng 327 ) 328 329 vocabulary.update_core_integrations( core_integrations ) 330 331 if ( np_rng.random() < config.underrepresented_integration_prob and 332 current_idx >= config.begin_underrepresented_integration_at 333 ): 334 user_prompt, underrepresented_integrations = add_underrepresented_integrations( 335 user_prompt, 336 vocabulary=vocabulary, 337 min_k=config.min_underrepresented_integrations, 338 max_k=config.max_underrepresented_integrations, 339 np_rng=np_rng 340 ) 341 342 # Note, on new file instances starts empty again 343 if ( np_rng.random() < config.recursive_integration_prob and 344 current_idx >= config.begin_recursive_integration_at and 345 len( instances ) > 20 346 ): 347 348 user_prompt, recursive_integration = add_recursive_integrations( 349 user_prompt, 350 previous_passages=[ instance[ 'completion' ] for instance in instances ], 351 np_rng=np_rng 352 ) 353 354 if np_rng.random() < config.moral_integration_prob : 355 user_prompt, moral_integrations = add_moral_integrations( 356 user_prompt, 357 vocabulary=vocabulary, 358 min_k=config.min_moral_integrations, 359 max_k=config.max_moral_integrations, 360 np_rng=np_rng 361 362 ) 363 vocabulary.update_moral_integrations( moral_integrations ) 364 365 user_prompt += f" {exposition_setting[ 'user prompt reminder' ]}" 366 367 batch.append( { 368 "explanitory_style": explanitory_style, 369 "concept_choice": concept_choice, 370 "core_integrations" : core_integrations, 371 "moral_integrations" : moral_integrations, 372 "recursive_integration" : recursive_integration, 373 "underrepresented_integrations" : underrepresented_integrations, 374 "prompt": user_prompt, 375 "index" : current_idx, 376 "refinements" : [] 377 } ) 378 379 if len( batch ) == batch_size or current_idx == end_idx - 1 : 380 381 start = time.time() 382 vsz_before = vocabulary.size() 383 n_tried = 0 384 385 batch_instances = await generate_batch( 386 semaphore=semaphore, 387 client=client, 388 model=model, 389 system_prompt=system_prompt, 390 refinement_system_prompt=exposition_setting[ 'refinement system prompt' ], 391 batch=batch, 392 vocabulary=vocabulary, 393 known_names=known_names, 394 refine=config.refine_vocabulary, 395 refine_at=config.begin_vocabulary_refinement_at, 396 retries_on_new=config.retries_on_new_word, 397 retry_on_new_at=config.retry_on_new_words_at, 398 retry_on_new_form_at=config.retry_on_new_word_forms_at, 399 rand_seed=rand_seed, 400 max_length=config.max_length, 401 min_length=config.min_length, 402 max_request_attempts=10 403 ) 404 405 for idx, c in enumerate( batch_instances ) : 406 407 instance = batch[ idx ] 408 409 text = c[0] 410 words = c[1] 411 refinements = c[2] 412 new_vocab_size = c[3] 413 n_tried += c[4] 414 415 instance[ 'completion' ] = text 416 instance[ 'vocab_size' ] = new_vocab_size 417 instance[ 'refinements'] = refinements 418 419 total_words += len( words ) 420 batch_words += len( words ) 421 422 # Update the integration trackers 423 424 instances.append( instance ) 425 426 if verbose : 427 print( f"-----------------------------------------\n\n{ORANGE}INSTANCE:{RESET} {instance['index']}\n" ) 428 print( f"{GREEN}PROMPT:{RESET} {instance['prompt']}\n" ) 429 for s, r in refinements : 430 print( f"{PINK}SENTENCE REFINEMENT:{RESET}\n{s}->\n{r}\n" ) 431 print( f"{YELLOW}COMPLETION:{RESET} {text}\n" ) 432 433 np_rng_state = np_rng.bit_generator.state 434 n_completed = current_idx+1 435 tc = time.time() - start 436 time_spent += tc 437 avg_time = time_spent / n_completed 438 avg_words = total_words / n_completed 439 avg_words_per_s = total_words / time_spent 440 vsz_now = vocabulary.size() 441 n_new_words = vsz_now - vsz_before 442 443 exp_boosts = ( 0 if current_idx < config.begin_underrepresented_integration_at 444 else len(batch) * config.underrepresented_integration_prob * ( config.max_underrepresented_integrations + config.min_underrepresented_integrations ) / 2.0 445 ) 446 447 print( f"retries: {n_tried-batch_size}" ) 448 print( f"batch w/s: {batch_words/tc:.2f}" ) 449 print( f"avg w/s: {avg_words_per_s:.2f}" ) 450 print( f"batch s/cmpl: {tc/batch_size:.3f}" ) 451 print( f"avg s/cmpl: {avg_time:.3f}" ) 452 print( f"avg w/cmpl: {avg_words:.2f}" ) 453 print( f"total c: {n_completed}" ) 454 print( f"total w: {total_words}, ~{round(total_words/0.75)}tk" ) 455 print( f"vocab size: {vsz_now}" ) 456 print( f"new words: {n_new_words}" ) 457 print( f"exp boosts: {exp_boosts}\n" ) 458 459 if ( batch_idx + 1 ) % save_every == 0 : 460 461 safe_save( { 462 'instances' : instances, 463 "config" : config, 464 "np_rng_state" : serialize_np_rng_state( dict(np_rng_state) ), 465 "random_seed" : rand_seed, 466 "time_spent" : time_spent, 467 'total_words' : total_words 468 }, filepath=output_path ) 469 470 safe_save( 471 { "np_rng_state" : serialize_np_rng_state( dict(np_rng_state) ) }, 472 filepath=random_state_path 473 ) 474 475 vocabulary.to_file( vocabulary_path ) 476 477 print( "Saved.\n" ) 478 479 batch = [] 480 batch_idx += 1 481 batch_words = 0 482 483 safe_save( { 484 'instances' : instances, 485 "config" : config, 486 "np_rng_state" : serialize_np_rng_state( dict(np_rng_state) ), 487 "random_seed" : rand_seed, 488 "time_spent" : time_spent, 489 'total_words' : total_words 490 }, filepath=output_path ) 491 492 safe_save( 493 { "np_rng_state" : serialize_np_rng_state( dict(np_rng_state) ) }, 494 filepath=random_state_path 495 ) 496 497 vocabulary.to_file( vocabulary_path )