accessible_worlds.aw_generate_morphology

View Source

  1import asyncio
  2import argparse
  3from pathlib import Path
  4from typing import Any
  5from openai import AsyncOpenAI
  6
  7from .aw_requests import get_response
  8from .aw_utils import in_dictionary, safe_save, load_json
  9from .aw_vocabulary import AwVocabulary
 10
 11async def get_morphological_splits(
 12    api_key : str,
 13    api_url : str,
 14    model : str,
 15    words : list[str],
 16    names : set[str],
 17    n_concurrent : int,
 18    batch_size : int,
 19) -> dict[str, dict[str,Any] ] :
 20
 21    """
 22    Note: results are only as good as the model, and Gemma 26b a4b produces insatisfactory intial results. Post-hoc refinement is possible. Otherwise, a better model, or alternative approach should be used.  
 23    """
 24
 25    max_tries = 1
 26
 27    semaphore = asyncio.Semaphore( n_concurrent )
 28
 29    client = AsyncOpenAI(
 30        base_url=api_url,
 31        api_key=api_key
 32    )
 33
 34    final_result = {}
 35
 36    prompt_template = "You are working through a list of words one at a time, and your task is to apply morphological splitting. For example, sunflower becomes sun-flower, deformed becomes de-form-ed, atomic words like gate would stay the same. Plural words split on the s, or es, e.g. cakes becomes cake-s, and outdoors becomes out-door-s, while foxes becomes fox-es. Some words will be plural names, which follow the same rule, e.g., Elaras produces Elara-s, Alices produces Alice-s. Output only the split word, with its components separated by dashes, except if the root word is an allomorph (altered by spelling rules), then output both the version surface-level string split and the lemmatized dictionary split separated by a pipe (|). For example, unverified becomes  un-verifi-ed | un-verify-ed, happiness becomes happi-ness | happy-ness, communicating becomes communicat-ing | communicate-ing, berries becomes berri-es | berry-es, continues becomes continue-s, injured becomes injur-ed | injure-ed, moving becomes mov-ing | move-ing, carries becomes carri-es | carry-es, worried becomes worri-ed | worry-ed, named becomes nam-ed | name-ed, and consonant doubling words like running become, e.g., runn-ing | run-ing. The surface level split (left side) with the dashes removed reproduce the original word exactly. Here is the next word: \\n\\nWORD"
 37
 38    async def generate_one( idx : int ) -> tuple[ int, list[str], bool ] :
 39
 40        word = words[ idx ]
 41
 42        seen = set()
 43
 44        if in_dictionary( word ) or ( word in names and word[-1] == 's' ) : 
 45
 46            prompt = prompt_template.replace( "WORD", word )
 47            is_valid = False
 48            result = []
 49            response = ""
 50
 51            error_message = ""
 52
 53            for i in range( max_tries ) :
 54
 55                user_prompt = prompt
 56
 57                if error_message : 
 58                    user_prompt += f"The previous attempt was {response}, which failed because: {error_message} Can you try again?"  
 59
 60                response = await get_response( 
 61                    semaphore=semaphore,
 62                    client=client,
 63                    model=model,
 64                    system_prompt="",
 65                    user_prompt=prompt+error_message,
 66                    seed=i,
 67                    temperature=0.0,
 68                    reasoning_budget=0,
 69                    max_request_attempts=10
 70                )
 71
 72                print( f"{word} > {response}" )
 73
 74                error_message = ""
 75
 76                splits = response.split( "|" )
 77                splits = [ split.strip() for split in splits ]
 78
 79                result = splits
 80
 81                # Valid if left side reproduces word dashes removed
 82                # and if there is a right side, it is aligned with the left side 
 83                # and where they differ, the right side's difering token must be a dictionary word
 84
 85                left_side_valid  = splits[ 0 ].replace( "-", "" ) == word
 86                right_side_valid = len( splits ) == 0
 87
 88                if not left_side_valid :
 89                    error_message += f"The left side failed to reproduce the {word} without the dashes. "
 90
 91                if len( splits ) > 1 :
 92
 93                    s1 = splits[0]
 94                    s2 = splits[1]
 95
 96                    s1 = [ s.strip() for s in s1.split( '-'  ) ]
 97                    s2 = [ s.strip() for s in s2.split( '-'  ) ]
 98
 99                    if s1 == s2 :
100                        result = [ splits[0] ]
101                    elif len( s1 ) != len( s2 ) :
102                        right_side_valid = False 
103                        error_message += "Left and right sides not aligned. "
104                    else :
105
106                        for si, _ in enumerate( s1 ) :
107                            if s1[ si ] != s2[ si ] :
108                                if not ( 
109                                    in_dictionary( s2[ si ] )
110                                ) :
111                                    right_side_valid = False
112                                    error_message += "The lemmatized component of the right side is not a dictionary word. "
113
114                is_valid = left_side_valid and right_side_valid
115
116                if is_valid :
117                    break
118        else :
119            result = [ word ] 
120            is_valid = True
121
122        return idx, result, is_valid
123
124    for start_idx in range( 0, len(words), batch_size ) :
125
126        end_idx = min( start_idx + batch_size, len( words ) )
127
128        async with asyncio.TaskGroup() as tg:
129            tasks = [
130                tg.create_task(generate_one(idx)) 
131                for idx in range( start_idx, end_idx )
132            ]
133        
134        batch_result = [task.result() for task in tasks]
135        for res in batch_result : 
136            final_result[ words[ res[0] ] ] = {
137                "splits"   : res[1]
138            }
139
140    return final_result
141
142if __name__ == "__main__":
143
144    parser = argparse.ArgumentParser( description="Perform generative morphological split." )
145
146    parser.add_argument( "--api_url",     type=str )
147    parser.add_argument( "--api_key",     type=str )
148    parser.add_argument( "--model",       type=str )
149    parser.add_argument( "--vocab_path",  type=str,              help="Path to vocabulary file" )
150    parser.add_argument( "--output_dir",  type=str,              help="Output directory" )
151    parser.add_argument( "--concurrency", type=int, default=1,   help="Number of concurrent requests (default: 1)" )
152    parser.add_argument( "--batch_size",  type=int, default=1,   help="Batch size (default: 1)" )
153
154    args = parser.parse_args()
155
156    api_url     = args.api_url
157    api_key     = args.api_key
158    model       = args.model
159    vocab_path  = Path( args.vocab_path )
160    output_dir  = Path( args.output_dir )
161    concurrency = args.concurrency
162    batch_size  = args.batch_size
163
164    if not vocab_path.exists() :
165        raise ValueError( f"Vocabulary file {vocab_path} doesn't exist." )
166
167    if not output_dir.exists() :
168        raise ValueError( f"Output directory {output_dir} doesn't exist." )
169
170    vocabulary = AwVocabulary.from_file( vocab_path )
171
172    words = vocabulary.get_words()
173    names = vocabulary.get_names()
174
175    print( f"{concurrency} {batch_size}" )
176
177    # results = asyncio.run( get_morphological_splits( 
178    #     api_key=api_key,
179    #     api_url=api_url,
180    #     model=model,
181    #     words=words,
182    #     names=names,
183    #     n_concurrent=concurrency,
184    #     batch_size=batch_size
185    # ) )
186
187    results = load_json( output_dir / "morphology.json"  )
188    fixes = load_json(  output_dir / "fixes.json" )
189
190    new_results = {
191        word : r["splits"] 
192        for word, r in results.items()
193    }
194
195    for fix in fixes :
196        if fix in results : 
197            new_results[ fix ] = fixes[ fix ][ "splits" ]
198
199    safe_save( data=new_results, filepath=( output_dir / "morphology-fixed.json" )  )

async def get_morphological_splits( api_key: str, api_url: str, model: str, words: list[str], names: set[str], n_concurrent: int, batch_size: int) -> dict[str, dict[str, typing.Any]]: View Source

 12async def get_morphological_splits(
 13    api_key : str,
 14    api_url : str,
 15    model : str,
 16    words : list[str],
 17    names : set[str],
 18    n_concurrent : int,
 19    batch_size : int,
 20) -> dict[str, dict[str,Any] ] :
 21
 22    """
 23    Note: results are only as good as the model, and Gemma 26b a4b produces insatisfactory intial results. Post-hoc refinement is possible. Otherwise, a better model, or alternative approach should be used.  
 24    """
 25
 26    max_tries = 1
 27
 28    semaphore = asyncio.Semaphore( n_concurrent )
 29
 30    client = AsyncOpenAI(
 31        base_url=api_url,
 32        api_key=api_key
 33    )
 34
 35    final_result = {}
 36
 37    prompt_template = "You are working through a list of words one at a time, and your task is to apply morphological splitting. For example, sunflower becomes sun-flower, deformed becomes de-form-ed, atomic words like gate would stay the same. Plural words split on the s, or es, e.g. cakes becomes cake-s, and outdoors becomes out-door-s, while foxes becomes fox-es. Some words will be plural names, which follow the same rule, e.g., Elaras produces Elara-s, Alices produces Alice-s. Output only the split word, with its components separated by dashes, except if the root word is an allomorph (altered by spelling rules), then output both the version surface-level string split and the lemmatized dictionary split separated by a pipe (|). For example, unverified becomes  un-verifi-ed | un-verify-ed, happiness becomes happi-ness | happy-ness, communicating becomes communicat-ing | communicate-ing, berries becomes berri-es | berry-es, continues becomes continue-s, injured becomes injur-ed | injure-ed, moving becomes mov-ing | move-ing, carries becomes carri-es | carry-es, worried becomes worri-ed | worry-ed, named becomes nam-ed | name-ed, and consonant doubling words like running become, e.g., runn-ing | run-ing. The surface level split (left side) with the dashes removed reproduce the original word exactly. Here is the next word: \\n\\nWORD"
 38
 39    async def generate_one( idx : int ) -> tuple[ int, list[str], bool ] :
 40
 41        word = words[ idx ]
 42
 43        seen = set()
 44
 45        if in_dictionary( word ) or ( word in names and word[-1] == 's' ) : 
 46
 47            prompt = prompt_template.replace( "WORD", word )
 48            is_valid = False
 49            result = []
 50            response = ""
 51
 52            error_message = ""
 53
 54            for i in range( max_tries ) :
 55
 56                user_prompt = prompt
 57
 58                if error_message : 
 59                    user_prompt += f"The previous attempt was {response}, which failed because: {error_message} Can you try again?"  
 60
 61                response = await get_response( 
 62                    semaphore=semaphore,
 63                    client=client,
 64                    model=model,
 65                    system_prompt="",
 66                    user_prompt=prompt+error_message,
 67                    seed=i,
 68                    temperature=0.0,
 69                    reasoning_budget=0,
 70                    max_request_attempts=10
 71                )
 72
 73                print( f"{word} > {response}" )
 74
 75                error_message = ""
 76
 77                splits = response.split( "|" )
 78                splits = [ split.strip() for split in splits ]
 79
 80                result = splits
 81
 82                # Valid if left side reproduces word dashes removed
 83                # and if there is a right side, it is aligned with the left side 
 84                # and where they differ, the right side's difering token must be a dictionary word
 85
 86                left_side_valid  = splits[ 0 ].replace( "-", "" ) == word
 87                right_side_valid = len( splits ) == 0
 88
 89                if not left_side_valid :
 90                    error_message += f"The left side failed to reproduce the {word} without the dashes. "
 91
 92                if len( splits ) > 1 :
 93
 94                    s1 = splits[0]
 95                    s2 = splits[1]
 96
 97                    s1 = [ s.strip() for s in s1.split( '-'  ) ]
 98                    s2 = [ s.strip() for s in s2.split( '-'  ) ]
 99
100                    if s1 == s2 :
101                        result = [ splits[0] ]
102                    elif len( s1 ) != len( s2 ) :
103                        right_side_valid = False 
104                        error_message += "Left and right sides not aligned. "
105                    else :
106
107                        for si, _ in enumerate( s1 ) :
108                            if s1[ si ] != s2[ si ] :
109                                if not ( 
110                                    in_dictionary( s2[ si ] )
111                                ) :
112                                    right_side_valid = False
113                                    error_message += "The lemmatized component of the right side is not a dictionary word. "
114
115                is_valid = left_side_valid and right_side_valid
116
117                if is_valid :
118                    break
119        else :
120            result = [ word ] 
121            is_valid = True
122
123        return idx, result, is_valid
124
125    for start_idx in range( 0, len(words), batch_size ) :
126
127        end_idx = min( start_idx + batch_size, len( words ) )
128
129        async with asyncio.TaskGroup() as tg:
130            tasks = [
131                tg.create_task(generate_one(idx)) 
132                for idx in range( start_idx, end_idx )
133            ]
134        
135        batch_result = [task.result() for task in tasks]
136        for res in batch_result : 
137            final_result[ words[ res[0] ] ] = {
138                "splits"   : res[1]
139            }
140
141    return final_result

Note: results are only as good as the model, and Gemma 26b a4b produces insatisfactory intial results. Post-hoc refinement is possible. Otherwise, a better model, or alternative approach should be used.