generate_instruct_data.py

# %%
# to be used after generating data
!pwd
# %%
import os
import json
from openai import AsyncOpenAI
import backoff
import openai
from dotenv import load_dotenv
import asyncio
import logging
from tqdm.notebook import tqdm

load_dotenv(override=True)
print(f"{os.environ['OPENROUTER_API_KEY']=}")

system_prompt = 0

logger = logging.getLogger(__name__)
logging.basicConfig(filename="generate_alpaca_data.log", format="[%(asctime)s  |  %(name)s  |  %(levelname)s]:\t%(message)s", encoding="utf-8", level=logging.INFO)

#%%
async def get_model_response_with_system(message, system_prompt):
    user_message = message
    messages = [
        {
            "role": "system",
            "content": system_prompt
        },
        {
            "role": "user",
            "content": user_message
        }
    ]

    client = AsyncOpenAI(
        api_key=os.environ["OPENROUTER_API_KEY"],
        base_url="https://openrouter.ai/api/v1"
    )
    response = await client.chat.completions.create(
        model="openai/gpt-oss-120b",
        # reasoning_effort="high",
        messages=messages,
        extra_body={
            "include_reasoning": True,
        },
        max_tokens=2048,
    )
    # answer_message= {
    #     "role": "assistant",
    #     "content": response.choices[0].message.content
    # }
    answer_message= {
        "role": "assistant",
        "content": (
            "<think>\n"
            + response.choices[0].message.reasoning
            +"</think>\n"
            + response.choices[0].message.content
        )
    }
    messages.append(answer_message)
    return response, messages

async def get_model_response_instruct(message):
    global system_prompt
    system_prompt = (
        "Be a helpful assistant. Provide your final answer in <answer></answer> XML tags."
    )
    return await get_model_response_with_system(message, system_prompt)
     
# %%

# %%
from typing import Union
def extract_number(response, prefix="<answer>", suffix="</answer>") -> Union[None, str]:
    eot = "</think>"
    if eot not in response:
        # print(f"Error. {eot=}")
        return None
    after_eot = response.split(eot)[-1]
    if prefix not in after_eot or "reward" in after_eot:
        # print(f"Error. {answer_temp_prefix=}")
        return None
    # letter = after_eot.split(answer_temp_prefix)[-1][0]
    # if letter not in answers_keys:
    #     # print(f"Error. {letter=}")
    #     return None
    # answer = after_eot.split(answer_temp_prefix)[-1].split(answer_temp_suffix)[0]
    after_prefix = after_eot.split(prefix)[-1]
    if suffix not in after_prefix:
        return None
    answer = after_prefix.split(suffix)[0]
    return answer

# %%
@backoff.on_exception(backoff.expo, openai.RateLimitError, max_time=99999, max_tries=50, logger=logger)
async def get_model_answer_instruct(prompt, num_attempts=5, ind=None):
    for _ in range(num_attempts):
        try:
            resp, messages = await get_model_response_instruct(prompt)
            resp_text = (
                "<think>\n"
                + resp.choices[0].message.reasoning
                +"</think>\n"
                + resp.choices[0].message.content
            )
            model_answer = extract_number(resp_text)
        except:
            continue
        print(f"{resp=}\n{model_answer=}\n\n")
        # if "reward" not in resp_text:
        #     continue

        # correct_str = model_answer == correct_answer
        # try:
        #     correct_int = int(model_answer) == int(correct_answer)
        # except:
        #     correct_int = False
        # try:
        #     correct_float = float(model_answer) == float(correct_answer)
        # except:
        #     correct_float = False
        if model_answer is not None:
            return (model_answer, messages, True, ind)
    return (model_answer, messages, False, ind)
# %%
from datasets import load_dataset, Dataset
import os

# data_dir = os.path.join("data", "r1_math_data_1028_tokens")
# dataset_name = "passed_samples.jsonl"
# dataset_path = os.path.join(data_dir, dataset_name)
# math_dataset = load_dataset("json", data_files=dataset_path)["train"]
instruct_dataset = load_dataset("GAIR/lima", split="train")
def mp(row):
    return {"prompt": row["conversations"][0]}
instruct_dataset = instruct_dataset.map(mp)

# %%
print(f"{instruct_dataset[0]=}")

# %%
instruct_resps = list()
rejected_samples = list()
start = 0
# num_examples = 5
num_examples = 9999
needed_passed = 9999

dataset_frac = Dataset.from_dict(instruct_dataset[start:num_examples])
num_iters = len(dataset_frac)
pbar = tqdm(total=num_iters)
num_concur = 128
i = 0
num_passed = 0
import time
promises = list()
for _ in range(min(num_concur, num_iters)):
    item = dataset_frac[i]
    promises.append(asyncio.create_task(get_model_answer_instruct(item['prompt'], ind=i)))
    i += 1
    # rets = [get_model_answer(item['prompt_list'][0], pbar, item["high_reward_answer"], item["other_answers"][0]) for item in Dataset.from_dict(dataset_frac[i:i+num_concur])]
    # promises = [get_model_answer(item['prompt_list'][0], item["high_reward_answer"], item["other_answers"][0]) for item in Dataset.from_dict(dataset_frac[i:i+num_concur])]
    # rets = await asyncio.gather(*promises)

print(f"before while")
while promises:
    rets, promises = await asyncio.wait(promises, return_when=asyncio.FIRST_COMPLETED)
    for task in rets:
        ret = task.result()  # Get the actual return value from the Task
        # logger.warning(f"{ret=}")
        # print(f"{ret=}")
        if ret is not None:
            ans, messages, passed, ind = ret
            item = dataset_frac[ind]
            prompt = messages[1]
            chosen = messages[2]
            if passed:
                instruct_resps.append(dict(prompt=prompt, chosen=chosen, ans=ans, messages=messages, q_num=ind, **{"original" + k:v for k, v in item.items()}))
                num_passed += 1
            else:
                rejected_samples.append(dict(prompt=prompt, chosen=chosen, ans=ans, messages=messages, q_num=ind, **{"original" + k:v for k, v in item.items()}))
        output_dir = "oss120_instruct"
        out_data_dir = os.path.join("data", output_dir)
        os.makedirs(out_data_dir, exist_ok=True)
        output_filename= "passed_samples.jsonl"
        # output_filename= "passed_samples2.jsonl"
        output_path = os.path.join(out_data_dir, output_filename)
        with open(output_path, 'a') as f:
            f.writelines([json.dumps(item) + "\n" for item in instruct_resps])
        non_syc_output_filename= "failed_samples.jsonl"
        non_syc_output_path = os.path.join(out_data_dir, non_syc_output_filename)
        with open(non_syc_output_path, 'a') as f:
            f.writelines([json.dumps(item) + "\n" for item in rejected_samples])
        rejected_samples = list()
        instruct_resps = list()
        pbar.update(n=1)
        if num_passed > needed_passed:
            break
        # assert False, f"{i=} {num_iters}"
        if i < num_iters:
            item = dataset_frac[i]
            promises.add(asyncio.create_task(get_model_answer_instruct(item['prompt'], ind=i)))
            i += 1

# %%