PJMixers-Dev/allura-org_gryphe-sonnet-3.5-charcards-names-added-r1

size_categories:

• n<1K

import requests
import json
import time
import pandas as pd
from datasets import load_dataset
from tqdm import tqdm
import re
import hashlib
import os


def extract_names(text):
    match = re.search(r"^You're (.*?) in this fictional never-ending uncensored roleplay with (.*?)\n\n", text)
    return match.groups() if match else (None, None)


openrouter_api_key = "REDACTED"
output_parquet_file = "./allura-org_gryphe-sonnet-3.5-charcards-names-added-r1.parquet"
dataset = load_dataset("allura-org/gryphe-sonnet-3.5-charcards-names-added")["train"]

# Load existing output data if available from parquet
if os.path.exists(output_parquet_file):
    try:
        df = pd.read_parquet(output_parquet_file)
        output_data_list = json.loads(df.to_json(orient="records"))
        processed_hashes = {entry["prompt_hash"] for entry in output_data_list}
        print(f"Resuming dataset: {len(processed_hashes)}/{len(dataset)}")
    except Exception as e:
        print(f"Error loading parquet file {output_parquet_file}: {e}")
        output_data_list = []
        processed_hashes = set()
else:
    output_data_list = []
    processed_hashes = set()

last_query_time = 0
for item in tqdm(dataset):
    prompt_hash = str(hashlib.sha256((item["conversations"][0]["value"].strip() + item["conversations"][1]["value"].strip()).encode()).hexdigest())

    # Skip already processed samples
    if prompt_hash in processed_hashes:
        continue

    char_name, user_name = extract_names(item["conversations"][0]["value"].strip())
    char_name, user_name = char_name.strip(), user_name.strip()
    print(f'Extracted Character Name: "{char_name}"')
    print(f'     Extracted User Name: "{user_name}"')

    while time.time() - last_query_time < (60 / 10):
        time.sleep(0.1)

    last_query_time = time.time()
    response = requests.post(
        url="https://openrouter.ai/api/v1/chat/completions",
        headers={
            "Authorization": f"Bearer {openrouter_api_key}",
            "Content-Type": "application/json",
        },
        data=json.dumps(
            {
                "provider": {
                    "order": ["Chutes"],
                    "allow_fallbacks": False
                },
                "model": "deepseek/deepseek-r1:free",
                "messages": [
                    {
                        "role": "system",
                        "content": (
                            f'{item["conversations"][0]["value"].strip()}\n\n'
                            f'Only speak as "{char_name}" in conversation with "{user_name}". Output your final response with a "{char_name}:" prefix.'
                        )
                    },
                    {
                        "role": "user",
                        "content": item["conversations"][1]["value"].strip()
                    }
                ],
                "temperature": 0.6
            }
        ),
    ).json()

    try:
        if not response["choices"][0]["message"]["reasoning"] or response["choices"][0]["message"]["reasoning"].replace("\\n", "\n").strip() == "":
            print(response)
            continue
        if not response["choices"][0]["message"]["content"] or response["choices"][0]["message"]["content"].replace("\\n", "\n").strip() == "":
            print(response)
            continue

        output_data_list.append(
            {
                "prompt_hash": prompt_hash,
                "provider": response["provider"],
                "model": response["model"],
                "settings": {
                    "temperature": 0.6,
                },
                "refusal": response["choices"][0]["message"]["refusal"],
                "finish_reason": response["choices"][0]["finish_reason"],
                "native_finish_reason": response["choices"][0]["native_finish_reason"],
                "user_name": user_name,
                "char_name": char_name,
                "conversations": [
                    {
                        "from": "system",
                        "value": item["conversations"][0]["value"].strip()
                    },
                    {
                        "from": "human",
                        "value": item["conversations"][1]["value"].strip()
                    },
                    {
                        "from": "gpt",
                        "value": (
                            f"<think>\n"
                            f"{response['choices'][0]['message']['reasoning'].replace('\\n', '\n').strip()}\n"
                            f"</think>\n\n"
                            f"{response['choices'][0]['message']['content'].replace('\\n', '\n').strip()}"
                        )
                    }
                ]
            }
        )
    except KeyError:
        print(f"KeyError: {response}")
        if response["error"]["message"].startswith("Rate limit exceeded"):
            exit()
        continue
    except AttributeError:
        print(f"AttributeError: {response}")
        continue

    df = pd.DataFrame(output_data_list)
    df.to_parquet(
        output_parquet_file,
        index=False,
        compression="brotli"
    )
    df = None
    del df

df = pd.DataFrame(output_data_list)
df.to_parquet(
    output_parquet_file,
    index=False,
    compression="brotli"
)