import asyncio import sys from typing import AsyncContextManager, Dict, List, Callable, Any import logging from random import sample from thrift.py3.client import get_client from thrift.py3.common import Protocol from picard.clients import Picard from picard.types import ( FeedTimeoutFailure, SQLSchema, RegisterSQLSchemaException, FeedException, Mode, ColumnType, ) from seq2seq.utils.picard_model_wrapper import PicardLauncher try: from transformers.models.auto import AutoTokenizer from transformers.tokenization_utils_fast import PreTrainedTokenizerFast from alive_progress import alive_bar has_tokenizer = True except: PreTrainedTokenizerFast = Any has_tokenizer = False logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", handlers=[logging.StreamHandler(sys.stdout)], level=logging.WARNING, ) logger = logging.getLogger(__name__) async def test_register_sql_schema( sql_schemas: Dict[str, SQLSchema], get_picard_client: Callable[[], AsyncContextManager[Picard]], ) -> None: async with get_picard_client() as client: for db_id, sql_schema in sql_schemas.items(): logger.info(f"registering {db_id}...") try: await client.registerSQLSchema(db_id, sql_schema) except RegisterSQLSchemaException: # db already registered logger.warning(f"schema already registered: {db_id}") pass logger.info(f"{db_id} registered") async def test_register_tokenizer( tokenizer: PreTrainedTokenizerFast, get_picard_client: Callable[[], AsyncContextManager[Picard]], ) -> None: async with get_picard_client() as client: logger.info(f"registering {tokenizer}...") json_str = tokenizer.backend_tokenizer.to_str(pretty=False) await client.registerTokenizer(json_str) logger.info(f"{tokenizer} registered") async def _feed(client: Picard, input_ids: List[int], token: int, with_guards: bool = True) -> None: try: await client.feed( input_ids, token, Mode.PARSING_WITH_GUARDS if with_guards else Mode.PARSING_WITHOUT_GUARDS, ) except FeedException: pass async def test_feedParser( input_ids: List[List[int]], tokens: List[int], get_picard_client: Callable[[], AsyncContextManager[Picard]], with_guards: bool = True, ) -> None: async with get_picard_client() as client: logger.info( f"synchronously feeding {len(tokens)} token continuations for {len(input_ids)} input sequence(s)..." ) with alive_bar(len(tokens) * len(input_ids)) as bar: for input_ids_ in input_ids: for token in tokens: await _feed(client, input_ids_, token, with_guards) bar() logger.info(f"synchronously feeding complete") async with get_picard_client() as client: logger.info( f"asynchronously feeding {len(tokens)} token continuations for {len(input_ids)} input sequence(s)..." ) futures = [_feed(client, input_ids_, token, with_guards) for token in tokens for input_ids_ in input_ids] with alive_bar(len(tokens) * len(input_ids)) as bar: for f in asyncio.as_completed(futures): await f bar() logger.info(f"asynchronously feeding complete") def in_vitro_picard( input_ids: List[List[int]], tokens: List[int], max_tokens_to_check: int, get_picard_client: Callable[[], AsyncContextManager[Picard]], with_guards: bool = True, batched: bool = True, ): async def _simulate_step(input_ids: List[List[int]], top_tokens: List[List[int]]): async with get_picard_client() as client: futures = [ # asyncio.sleep(0.01) _feed( client=client, input_ids=input_ids_batch, token=top_token, with_guards=with_guards, ) for _batch_idx, (input_ids_batch, top_token_batch) in enumerate(zip(input_ids, top_tokens)) for top_token in top_token_batch ] for f in asyncio.as_completed(futures): await f async def _simulate_batched_step(input_ids: List[List[int]], top_tokens: List[List[int]]): async with get_picard_client() as client: try: results = await client.batchFeed( input_ids, top_tokens, Mode.PARSING_WITH_GUARDS if with_guards else Mode.PARSING_WITHOUT_GUARDS, ) # for r in results: # if isinstance(r.feedResult.value, FeedTimeoutFailure): # raise RuntimeError except FeedException: pass beam_size = len(input_ids) assert beam_size > 0 seq_len = len(input_ids[0]) assert all(len(seq) == seq_len for seq in input_ids) assert len(tokens) >= max_tokens_to_check - 1 for n in range(seq_len): truncated_input_ids = [seq[:n] for seq in input_ids] simulated_top_tokens = [ [seq[n]] + sample(tokens, k=len(tokens))[: max_tokens_to_check - 1] for seq in input_ids ] asyncio.run( _simulate_batched_step(input_ids=truncated_input_ids, top_tokens=simulated_top_tokens) if batched else _simulate_step(input_ids=truncated_input_ids, top_tokens=simulated_top_tokens) ) sql_schemas = { "test": SQLSchema( columnNames={"0": "column"}, columnTypes={"0": ColumnType.NUMBER}, tableNames={"0": "table"}, columnToTable={"0": "0"}, tableToColumns={"0": ["0"]}, foreignKeys=dict(), primaryKeys=list(), ), "car_1": SQLSchema( columnNames={ "1": "ContId", "10": "ModelId", "11": "Maker", "12": "Model", "13": "MakeId", "14": "Model", "15": "Make", "16": "Id", "17": "MPG", "18": "Cylinders", "19": "Edispl", "2": "Continent", "20": "Horsepower", "21": "Weight", "22": "Accelerate", "23": "Year", "3": "CountryId", "4": "CountryName", "5": "Continent", "6": "Id", "7": "Maker", "8": "FullName", "9": "Country", }, columnTypes={ "1": ColumnType.NUMBER, "10": ColumnType.NUMBER, "11": ColumnType.NUMBER, "12": ColumnType.TEXT, "13": ColumnType.NUMBER, "14": ColumnType.TEXT, "15": ColumnType.TEXT, "16": ColumnType.NUMBER, "17": ColumnType.TEXT, "18": ColumnType.NUMBER, "19": ColumnType.NUMBER, "2": ColumnType.TEXT, "20": ColumnType.TEXT, "21": ColumnType.NUMBER, "22": ColumnType.NUMBER, "23": ColumnType.NUMBER, "3": ColumnType.NUMBER, "4": ColumnType.TEXT, "5": ColumnType.NUMBER, "6": ColumnType.NUMBER, "7": ColumnType.TEXT, "8": ColumnType.TEXT, "9": ColumnType.TEXT, }, tableNames={ "0": "continents", "1": "countries", "2": "car_makers", "3": "model_list", "4": "car_names", "5": "cars_data", }, columnToTable={ "1": "0", "10": "3", "11": "3", "12": "3", "13": "4", "14": "4", "15": "4", "16": "5", "17": "5", "18": "5", "19": "5", "2": "0", "20": "5", "21": "5", "22": "5", "23": "5", "3": "1", "4": "1", "5": "1", "6": "2", "7": "2", "8": "2", "9": "2", }, tableToColumns={ "0": ["1", "2"], "1": ["3", "4", "5"], "2": ["6", "7", "8", "9"], "3": ["10", "11", "12"], "4": ["13", "14", "15"], "5": ["16", "17", "18", "19", "20", "21", "22", "23"], }, foreignKeys={"11": "6", "14": "12", "16": "13", "5": "1", "9": "3"}, primaryKeys=["1", "3", "6", "10", "13", "16"], ), } def main() -> None: def get_picard_client() -> AsyncContextManager[Picard]: return get_client(Picard, host="127.0.0.1", port=9090, timeout=1, protocol=Protocol.BINARY) asyncio.run(test_register_sql_schema(sql_schemas=sql_schemas, get_picard_client=get_picard_client)) if has_tokenizer: tokenizer = AutoTokenizer.from_pretrained("t5-base", use_fast=True) asyncio.run(test_register_tokenizer(tokenizer=tokenizer, get_picard_client=get_picard_client)) tokens = list(tokenizer.get_vocab().values()) # they'll come in a random order input_ids = tokenizer( [ "<pad>test | select table.column from table</s>", "<pad>car_1 | select avg(edispl) from cars_data where id in (select id from cars_data where id in (select id from cars_data where id in (select id from cars_data where id in (select id from cars_data where id in (select id from cars_data where id in (select id from cars_data where", '<pad>car_1 | select t1.countryid, t1.countryname from countries as t1 join car_makers as t2 on t1.countryid = t2.country group by t1.countryid having count(*) > 3 union select t1.countryid, t1.countryname from countries as t1 join car_makers as t2 on t1.countryid = t2.country join model_list as t3 on t2.id = t3.maker where t3.model = "Fiat"</s>', ], add_special_tokens=False, )["input_ids"] # fmt: off assert input_ids == [ [0, 794, 1820, 1738, 953, 5, 3297, 440, 29, 45, 953, 1], [0, 443, 834, 536, 1820, 1738, 3, 9, 208, 122, 599, 15, 10475, 40, 61, 45, 2948, 834, 6757, 213, 3, 23, 26, 16, 41, 7, 15, 3437, 3, 23, 26, 45, 2948, 834, 6757, 213, 3, 23, 26, 16, 41, 7, 15, 3437, 3, 23, 26, 45, 2948, 834, 6757, 213, 3, 23, 26, 16, 41, 7, 15, 3437, 3, 23, 26, 45, 2948, 834, 6757, 213, 3, 23, 26, 16, 41, 7, 15, 3437, 3, 23, 26, 45, 2948, 834, 6757, 213, 3, 23, 26, 16, 41, 7, 15, 3437, 3, 23, 26, 45, 2948, 834, 6757, 213, 3, 23, 26, 16, 41, 7, 15, 3437, 3, 23, 26, 45, 2948, 834, 6757, 213], [0, 443, 834, 536, 1820, 1738, 3, 17, 5411, 17529, 23, 26, 6, 3, 17, 5411, 17529, 4350, 45, 1440, 38, 3, 17, 536, 1715, 443, 834, 8910, 38, 3, 17, 357, 30, 3, 17, 5411, 17529, 23, 26, 3274, 3, 17, 4416, 17529, 563, 57, 3, 17, 5411, 17529, 23, 26, 578, 3476, 599, 1935, 61, 2490, 220, 7021, 1738, 3, 17, 5411, 17529, 23, 26, 6, 3, 17, 5411, 17529, 4350, 45, 1440, 38, 3, 17, 536, 1715, 443, 834, 8910, 38, 3, 17, 357, 30, 3, 17, 5411, 17529, 23, 26, 3274, 3, 17, 4416, 17529, 1715, 825, 834, 3350, 38, 3, 17, 519, 30, 3, 17, 4416, 23, 26, 3274, 3, 17, 5787, 8337, 213, 3, 17, 5787, 21770, 3274, 96, 3183, 144, 121, 1], ] # fmt: on # In-vitro Picard trials = 100 max_tokens_to_check = 2 simulated_beam_size = 4 with alive_bar(trials * max_tokens_to_check * simulated_beam_size) as bar: for _ in range(trials): in_vitro_picard( input_ids=[input_ids[2]] * simulated_beam_size, tokens=tokens[: max_tokens_to_check * 10], max_tokens_to_check=max_tokens_to_check, get_picard_client=get_picard_client, with_guards=True, ) bar(max_tokens_to_check * simulated_beam_size) # Populate parser cache # asyncio.run( # test_feedParser( # input_ids=[ # input_ids[2][:103], # ], # tokens=[input_ids[2][103]], # get_picard_client=get_picard_client, # with_guards=True, # ) # ) # Test performance with input id sequence of intermediate length # asyncio.run( # test_feedParser( # input_ids=[ # input_ids[2][:104], # ], # tokens=tokens[:1000], # get_picard_client=get_picard_client, # with_guards=True, # ) # ) # Test scaling with short input id sequence and many tokens # asyncio.run( # test_feedParser( # input_ids=[ # input_ids[0][:-1], # ], # do not add </s> so that the sequence is not terminated # tokens=tokens, # get_picard_client=get_picard_client, # with_guards=True, # ) # ) # Test timeout with adversarial input id sequence # asyncio.run( # test_feedParser( # input_ids=[input_ids[1] + [2029]], # tokens=tokens[:20], # get_picard_client=get_picard_client, # with_guards=True, # ) # ) if __name__ == "__main__": with PicardLauncher(): main()