feat: unified args

dataset_maker.py

@@ -103,12 +103,13 @@ prompts = {
     "wh": f"its semantic role",
 }
 
-async def classify_tokens(prompt: str, labels: dict[str, str], tokens: list[str], model="gpt-4o"):
+async def classify_tokens(args, prompt: str, labels: dict[str, str], tokens: list[str],
+                          model="gpt-4o"):
     tok_len = len(tokens)
     example = "[" + (", ".join([f'"{tok}"' for tok in tokens])) + "]"
     try:
         response = await client.chat.completions.create(
-            model=model, timeout=30,
+            model=args.openai_model, timeout=30,
             **({"reasoning_effort": "low"} if model.startswith("o") else {"presence_penalty": 0, "temperature": 0}),
             messages=[
                 {
@@ -166,20 +167,21 @@ async def classify_tokens(prompt: str, labels: dict[str, str], tokens: list[str]
         raise
 
 
-async def classify_with_retry(prompt, labels, tokens, model="gpt-4o", retry=10):
+async def classify_with_retry(args, prompt, labels, tokens, retry=10):
     for i in range(retry):
         try:
-            return await classify_tokens(prompt, labels, tokens, model=model)
+            return await classify_tokens(args, prompt, labels, tokens)
         except Exception as e:
             logger.error(f"attempt {i} failed {tokens} {prompt} {format_exc()}")
             await asyncio.sleep(i)
 
-async def generate_token_labels(case, model="gpt-4o"):
+
+async def generate_token_labels(args, case):
     tokens = sp_tokenize(case)
     sorted_cols = list(sorted(features.keys()))
     example = {}
     for idx, labels in  enumerate(list(await asyncio.gather(
-            *[classify_with_retry(prompts[col], features[col], tokens, model=model) for col in sorted_cols]))):
+            *[classify_with_retry(args, prompts[col], features[col], tokens) for col in sorted_cols]))):
         example[sorted_cols[idx]] = labels
     return example
 
@@ -229,7 +231,7 @@ async def main(args, cases):
         while len([t for t in tasks if t is not None]) >= max_concurrent_tasks:
             await asyncio.sleep(1)
         logger.info(f"scheduling case {case}")
-        tasks.append(asyncio.create_task(generate_token_labels(case, model=args.openai_model)))
+        tasks.append(asyncio.create_task(generate_token_labels(args, case)))
 
     # Block until done
     while len([t for t in tasks if t is not None]) > 0: