{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "\n",
    "tqdm.pandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Load raw data and convert to training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gzip\n",
    "import json\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "FILE_PATHS = [\n",
    "    'YOURPATH-no-hint-train-t05-run_1/output.with_completions.jsonl.gz',\n",
    "    'YOURPATH-no-hint-train-t05-run_2/output.with_completions.jsonl.gz',\n",
    "]\n",
    "\n",
    "# More memory efficient for large files\n",
    "# Initialize lists to store the data\n",
    "data = []\n",
    "\n",
    "\n",
    "# Read file line by line\n",
    "for FILE_PATH in FILE_PATHS:\n",
    "    with gzip.open(FILE_PATH, 'rb') as f:  # Use 'rb' for gzipped files\n",
    "        for i, line in tqdm(\n",
    "            enumerate(f), desc=f'Processing {FILE_PATH.split(\"/\")[-1]}'\n",
    "        ):\n",
    "            # Parse only the fields we need\n",
    "            raw_data = json.loads(line)\n",
    "            data.append(\n",
    "                {\n",
    "                    'resolved': raw_data['report']['resolved'],\n",
    "                    'messages': raw_data['raw_completions']['messages']\n",
    "                    if raw_data['raw_completions'] is not None\n",
    "                    else None,\n",
    "                    'git_patch': raw_data['test_result'].get('git_patch', ''),\n",
    "                    'tools': raw_data['raw_completions']['tools']\n",
    "                    if raw_data['raw_completions'] is not None\n",
    "                    and 'tools' in raw_data['raw_completions']\n",
    "                    else None,\n",
    "                }\n",
    "            )\n",
    "\n",
    "# Convert to DataFrame after collecting all data\n",
    "df = pd.DataFrame(data)\n",
    "print(f'#total amount of data={len(df)}')\n",
    "df = df[~df['messages'].isna()]\n",
    "print(f'#total amount of data after removing nan={len(df)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Filter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _contains_multiple_tool_calls(messages: list[dict]) -> bool:\n",
    "    return any(\n",
    "        message.get('tool_calls') and len(message['tool_calls']) > 1\n",
    "        for message in messages\n",
    "    )\n",
    "\n",
    "\n",
    "df['contains_multiple_tool_calls'] = df['messages'].apply(_contains_multiple_tool_calls)\n",
    "display(df.groupby(['contains_multiple_tool_calls'])['resolved'].sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import copy\n",
    "\n",
    "# Convert function calling messages to non-function calling messages\n",
    "from openhands.llm.fn_call_converter import (\n",
    "    FunctionCallConversionError,\n",
    "    convert_fncall_messages_to_non_fncall_messages,\n",
    "    convert_from_multiple_tool_calls_to_single_tool_call_messages,\n",
    ")\n",
    "\n",
    "total_failed = 0\n",
    "\n",
    "\n",
    "def _convert_messages(messages: list[dict], tools: list[dict]) -> list[dict]:\n",
    "    global total_failed\n",
    "    message_copy = copy.deepcopy(messages)\n",
    "    for message in message_copy:\n",
    "        if message['content'] is None:\n",
    "            message['content'] = ''\n",
    "    try:\n",
    "        return convert_fncall_messages_to_non_fncall_messages(\n",
    "            message_copy, tools, add_in_context_learning_example=False\n",
    "        )\n",
    "    except FunctionCallConversionError:\n",
    "        total_failed += 1\n",
    "        # print(f'Failed to convert messages: {messages}\\nTools: {tools}')\n",
    "        # traceback.print_exc()\n",
    "        return None\n",
    "\n",
    "\n",
    "df['converted_messages'] = df.apply(\n",
    "    lambda row: convert_from_multiple_tool_calls_to_single_tool_call_messages(\n",
    "        row['messages'], ignore_final_tool_result=True\n",
    "    ),\n",
    "    axis=1,\n",
    ")\n",
    "df['nonfncall_messages'] = df.apply(\n",
    "    lambda row: _convert_messages(row['converted_messages'], row['tools']), axis=1\n",
    ")\n",
    "print('total nan', df['nonfncall_messages'].isna().sum())\n",
    "df = df[~df['nonfncall_messages'].isna()]\n",
    "print(df['nonfncall_messages'].iloc[0])\n",
    "\n",
    "print(f'Total failed: {total_failed}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pandarallel import pandarallel\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "os.environ['TOKENIZERS_PARALLELISM'] = 'false'\n",
    "pandarallel.initialize(progress_bar=True, verbose=1, nb_workers=16)\n",
    "tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen2.5-7B-Instruct')\n",
    "\n",
    "\n",
    "def clean_messages(messages):\n",
    "    clean = []\n",
    "    for msg in messages:\n",
    "        if not isinstance(msg, dict):\n",
    "            continue\n",
    "        role = msg.get('role')\n",
    "        content = msg.get('content')\n",
    "        if isinstance(content, str):\n",
    "            text = content\n",
    "        elif isinstance(content, dict):\n",
    "            text = content.get('text')\n",
    "        elif (\n",
    "            isinstance(content, list)\n",
    "            and len(content) == 1\n",
    "            and isinstance(content[0], dict)\n",
    "        ):\n",
    "            text = content[0].get('text')\n",
    "        else:\n",
    "            print(f'Format not accepted {content}')\n",
    "        clean.append({'role': role, 'content': text})\n",
    "    return clean\n",
    "\n",
    "\n",
    "# Step 1: Clean the messages\n",
    "df['nonfncall_messages'] = df['nonfncall_messages'].apply(clean_messages)\n",
    "\n",
    "# Step 2: Compute token count\n",
    "df['n_tokens'] = df['nonfncall_messages'].parallel_apply(\n",
    "    lambda x: len(tokenizer.apply_chat_template(x))\n",
    ")\n",
    "\n",
    "# print(df['nonfncall_messages'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'BEFORE: #total={len(df)}')\n",
    "df_selected = df[df['n_tokens'] < 131072]\n",
    "print(f'AFTER(truncated to 128k): #total={len(df_selected)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_selected['n_tokens'].describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ecdf of n_tokens\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "display(df.groupby(['resolved'])['n_tokens'].describe())\n",
    "sns.ecdfplot(x='n_tokens', data=df, hue='resolved')\n",
    "plt.show()\n",
    "\n",
    "print(f'#total={len(df)}')\n",
    "df_selected = df[df['n_tokens'] < 131072]\n",
    "print(f'#selected={len(df_selected)}')\n",
    "display(df_selected.groupby(['resolved'])['n_tokens'].describe())\n",
    "sns.ecdfplot(x='n_tokens', data=df_selected, hue='resolved')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_selected[~df_selected['resolved']]['n_tokens'].describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_selected['resolved'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_selected.groupby(['resolved'])['n_tokens'].describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Save Resolved Messages for SFT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Flatten messages and change format to {\"content\": \"\", \"role\": \"\"}\n",
    "df_selected[df_selected['resolved']][['nonfncall_messages']].rename(\n",
    "    columns={'nonfncall_messages': 'messages'}\n",
    ").to_json(\n",
    "    os.path.join(\n",
    "        'PATH_TO_FILE',\n",
    "        f'policy_traj_128k_swegym_{df_selected[\"resolved\"].value_counts()[True]}i.jsonl',\n",
    "    ),\n",
    "    lines=True,\n",
    "    orient='records',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
