accessible_worlds.aw_sentence_refinement

View Source

  1from typing import Any
  2
  3import asyncio
  4from openai import AsyncOpenAI
  5
  6from .aw_vocabulary import AwVocabulary
  7
  8from .aw_utils import (
  9    split_sentences,
 10    split_words,
 11    split_paragraphs,
 12    PINK,
 13    RESET
 14)    
 15
 16from .aw_build_prompt import assemble_word_replacement_prompt
 17    
 18from .aw_requests import get_response
 19
 20async def refine_sentences(
 21    semaphore,
 22    client : AsyncOpenAI, 
 23    model : str,
 24    system_prompt : str,
 25    completion : str,
 26    new_words : set[str],
 27    vocabulary : AwVocabulary,
 28    current_idx : int,
 29    rand_seed : int
 30) -> tuple[ str, list[tuple[str,str]]] :
 31
 32    paragraphs = split_paragraphs( completion )
 33    paragraph_sentences = []
 34    sentences_with_new_words = []
 35
 36    # Gather all sentences that we will attempt to refine
 37
 38    for p_idx, paragraph in enumerate( paragraphs ) : 
 39                
 40        sentences = split_sentences( paragraph )
 41        paragraph_sentences.append( sentences )
 42
 43        for s_idx, sentence in enumerate( sentences ) :
 44
 45            words = split_words( sentence )
 46            
 47            new_sentence_words = [ word for word in words if word in new_words ]
 48
 49            # new_words that might have seen alternatives
 50            alternative_map = {
 51                word: alts
 52                for word in new_sentence_words
 53                if ( alts := vocabulary.synonyms( word, context=sentence ) )
 54            }
 55        
 56            if alternative_map : 
 57                sentences_with_new_words.append( ( p_idx, s_idx, alternative_map ) )
 58
 59    async def apply_refine_sentences(  
 60        p_idx : int, 
 61        s_idx : int,
 62        alternative_map : dict[str,Any]
 63    ) -> tuple[str,str] | None : 
 64
 65        sentence = paragraph_sentences[ p_idx ][ s_idx ] 
 66
 67        a_prompt = assemble_word_replacement_prompt(  
 68            paragraph=paragraphs[p_idx],
 69            sentence=sentence,
 70            alternatives=alternative_map
 71        )
 72
 73        replacement_sentence = await get_response( 
 74            semaphore=semaphore,
 75            client=client,
 76            model=model,
 77            system_prompt=system_prompt,
 78            user_prompt=a_prompt,
 79            seed=( rand_seed + current_idx * 0x9e3779b9 ) % ( 2**32 ),
 80            temperature=1.0,
 81            reasoning_budget=0
 82        )
 83
 84        if replacement_sentence is not None :   
 85            
 86            replacement_sentence.strip( '"' )   
 87            if replacement_sentence != sentence :   
 88                
 89                print( f"{PINK}SENTENCE REFINEMENT:{RESET}\n{sentence}->\n{replacement_sentence}\n" )   
 90                
 91                n_rp_words = len( split_words( replacement_sentence ) )
 92                n_original_words = len( split_words( sentence ) )   
 93                
 94                # Guards against the rare case the model fails to output only the sentence
 95                if abs( n_rp_words - n_original_words  ) < 5 : 
 96                    paragraph_sentences[ p_idx ][ s_idx ] = replacement_sentence
 97                else : 
 98                    print( "Unexpected refined sentence length, rejecting." )
 99
100                return sentence, replacement_sentence
101
102        return None
103
104    refinements = []
105    for p_idx, s_idx, alternative_map in sentences_with_new_words :
106        rf = await apply_refine_sentences( p_idx, s_idx, alternative_map )
107        if rf :
108            refinements.append( rf )
109
110    # Reassemble
111
112    for p_idx, sentences in enumerate( paragraph_sentences ) :
113        paragraphs[ p_idx ] = " ".join( sentences )
114
115    return "\n\n".join( paragraphs ), refinements

async def refine_sentences( semaphore, client: openai.AsyncOpenAI, model: str, system_prompt: str, completion: str, new_words: set[str], vocabulary: accessible_worlds.aw_vocabulary.AwVocabulary, current_idx: int, rand_seed: int) -> tuple[str, list[tuple[str, str]]]: View Source

 21async def refine_sentences(
 22    semaphore,
 23    client : AsyncOpenAI, 
 24    model : str,
 25    system_prompt : str,
 26    completion : str,
 27    new_words : set[str],
 28    vocabulary : AwVocabulary,
 29    current_idx : int,
 30    rand_seed : int
 31) -> tuple[ str, list[tuple[str,str]]] :
 32
 33    paragraphs = split_paragraphs( completion )
 34    paragraph_sentences = []
 35    sentences_with_new_words = []
 36
 37    # Gather all sentences that we will attempt to refine
 38
 39    for p_idx, paragraph in enumerate( paragraphs ) : 
 40                
 41        sentences = split_sentences( paragraph )
 42        paragraph_sentences.append( sentences )
 43
 44        for s_idx, sentence in enumerate( sentences ) :
 45
 46            words = split_words( sentence )
 47            
 48            new_sentence_words = [ word for word in words if word in new_words ]
 49
 50            # new_words that might have seen alternatives
 51            alternative_map = {
 52                word: alts
 53                for word in new_sentence_words
 54                if ( alts := vocabulary.synonyms( word, context=sentence ) )
 55            }
 56        
 57            if alternative_map : 
 58                sentences_with_new_words.append( ( p_idx, s_idx, alternative_map ) )
 59
 60    async def apply_refine_sentences(  
 61        p_idx : int, 
 62        s_idx : int,
 63        alternative_map : dict[str,Any]
 64    ) -> tuple[str,str] | None : 
 65
 66        sentence = paragraph_sentences[ p_idx ][ s_idx ] 
 67
 68        a_prompt = assemble_word_replacement_prompt(  
 69            paragraph=paragraphs[p_idx],
 70            sentence=sentence,
 71            alternatives=alternative_map
 72        )
 73
 74        replacement_sentence = await get_response( 
 75            semaphore=semaphore,
 76            client=client,
 77            model=model,
 78            system_prompt=system_prompt,
 79            user_prompt=a_prompt,
 80            seed=( rand_seed + current_idx * 0x9e3779b9 ) % ( 2**32 ),
 81            temperature=1.0,
 82            reasoning_budget=0
 83        )
 84
 85        if replacement_sentence is not None :   
 86            
 87            replacement_sentence.strip( '"' )   
 88            if replacement_sentence != sentence :   
 89                
 90                print( f"{PINK}SENTENCE REFINEMENT:{RESET}\n{sentence}->\n{replacement_sentence}\n" )   
 91                
 92                n_rp_words = len( split_words( replacement_sentence ) )
 93                n_original_words = len( split_words( sentence ) )   
 94                
 95                # Guards against the rare case the model fails to output only the sentence
 96                if abs( n_rp_words - n_original_words  ) < 5 : 
 97                    paragraph_sentences[ p_idx ][ s_idx ] = replacement_sentence
 98                else : 
 99                    print( "Unexpected refined sentence length, rejecting." )
100
101                return sentence, replacement_sentence
102
103        return None
104
105    refinements = []
106    for p_idx, s_idx, alternative_map in sentences_with_new_words :
107        rf = await apply_refine_sentences( p_idx, s_idx, alternative_map )
108        if rf :
109            refinements.append( rf )
110
111    # Reassemble
112
113    for p_idx, sentences in enumerate( paragraph_sentences ) :
114        paragraphs[ p_idx ] = " ".join( sentences )
115
116    return "\n\n".join( paragraphs ), refinements