{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Text2Persona — RealPersonaChat で `model.pt` を作る\n",
    "\n",
    "**目的**: `Personality-Recognition-on-RealPersonaChat` の `MultiTaskModel`(LUKE-japanese-base + 5つの回帰ヘッド) を `nu-dialogue/real-persona-chat` データで学習し、`model.pt` を作って Drive に保存する。\n",
    "\n",
    "**前提**\n",
    "- Runtime → Change runtime type → **GPU** (T4 で動くが、Pro/A100 だと現実的な時間で終わる)\n",
    "- Google Drive にマウント可能\n",
    "\n",
    "**フロー**\n",
    "1. ランタイム & Drive 確認\n",
    "2. リポジトリ clone + 依存インストール\n",
    "3. RealPersonaChat を HF datasets から読み込み\n",
    "4. モノローグ形式 CSV に変換 (file, speaker_id, dialogue, n, e, o, a, c)\n",
    "5. 話者単位で train/valid/test に分割 (リーク防止)\n",
    "6. `train.py` で学習\n",
    "7. DataParallel プレフィックス除去 → Drive に保存\n",
    "8. テストデータで推論サニティチェック"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 0 — GPU 確認"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "print('CUDA available:', torch.cuda.is_available())\n",
    "if torch.cuda.is_available():\n",
    "    print('Device:', torch.cuda.get_device_name(0))\n",
    "    print('VRAM (GB):', round(torch.cuda.get_device_properties(0).total_memory / 1e9, 2))\n",
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1 — Drive をマウントして作業ディレクトリを用意"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "\n",
    "import os\n",
    "WORK_DIR = '/content/drive/MyDrive/Text2Persona'\n",
    "MODEL_DIR = f'{WORK_DIR}/model/v1'\n",
    "LOG_DIR   = f'{WORK_DIR}/log'\n",
    "os.makedirs(MODEL_DIR, exist_ok=True)\n",
    "os.makedirs(LOG_DIR, exist_ok=True)\n",
    "print('WORK_DIR :', WORK_DIR)\n",
    "print('MODEL_DIR:', MODEL_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2 — リポジトリ clone & 依存\n",
    "\n",
    "`torch_geometric` は GCN 系の対話モデル変種でしか使わないので、モノローグ学習では入れない (Colab で入れると重いし壊れやすい)。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.chdir('/content')\n",
    "if not os.path.exists('/content/repo'):\n",
    "    !git clone https://github.com/fuyahuii/Personality-Recognition-on-RealPersonaChat.git repo\n",
    "os.chdir('/content/repo')\n",
    "print('repo:', os.getcwd())\n",
    "\n",
    "# We bypass the `datasets` library and pull the raw RealPersonaChat JSON via huggingface_hub,\n",
    "# so no `datasets` pin is needed.\n",
    "!pip install -q \\\n",
    "    \"transformers==4.36.0\" \\\n",
    "    \"sentencepiece\" \\\n",
    "    \"huggingface_hub>=0.20\" \\\n",
    "    \"pandas\" \\\n",
    "    \"scikit-learn\" \\\n",
    "    \"scipy\" \\\n",
    "    \"tqdm\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `train.py` の `torch_geometric` import を回避するパッチ\n",
    "\n",
    "`train.py` は冒頭で `from torch_geometric.nn import ...` を import している。モノローグ(context=0) では実際は使わないので、import 行を try/except でガードする。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "p = '/content/repo/train.py'\n",
    "with open(p, 'r', encoding='utf-8') as f:\n",
    "    src = f.read()\n",
    "\n",
    "old = 'from torch_geometric.nn import GraphConv,GATv2Conv,RGCNConv'\n",
    "new = (\n",
    "    'try:\\n'\n",
    "    '    from torch_geometric.nn import GraphConv,GATv2Conv,RGCNConv\\n'\n",
    "    'except Exception:\\n'\n",
    "    '    GraphConv = GATv2Conv = RGCNConv = None  # not needed for monologue mode'\n",
    ")\n",
    "if old in src and 'try:\\n    from torch_geometric' not in src:\n",
    "    src = src.replace(old, new)\n",
    "    with open(p, 'w', encoding='utf-8') as f:\n",
    "        f.write(src)\n",
    "    print('patched torch_geometric import')\n",
    "else:\n",
    "    print('already patched or marker not found')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Step 3 — RealPersonaChat を読み込む(ローカルアップロード版)\n\n事前に `real-persona-chat/real_persona_chat/` フォルダ(`interlocutors.json` + `dialogues/*.json`)を Colab に置く。**13,000 ファイル超なので zip でアップロードしてから unzip するのが速い**。\n\n候補パス(最初に見つかったものを使う):\n- `/content/drive/MyDrive/Text2Persona/real-persona-chat/real_persona_chat`  ← Drive に置くのが楽\n- `/content/real-persona-chat/real_persona_chat`  ← セッション中だけ\n\n**おすすめ手順**:\n```bash\n# ローカル Mac で\ncd \"/Users/.../Text2Persona/real-persona-chat\"\nzip -r real_persona_chat.zip real_persona_chat\n# Drive にアップロード → Colab で:\n!cp /content/drive/MyDrive/real_persona_chat.zip /content/\n!cd /content && unzip -q real_persona_chat.zip\n```"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "import json, os, glob\n\n# Look for the dataset in common upload locations.\nCANDIDATES = [\n    '/content/drive/MyDrive/Text2Persona/real-persona-chat/real_persona_chat',\n    '/content/drive/MyDrive/Text2Persona/real_persona_chat',\n    '/content/drive/MyDrive/real_persona_chat',\n    '/content/real-persona-chat/real_persona_chat',\n    '/content/real_persona_chat',\n]\n\ndata_root = None\nfor c in CANDIDATES:\n    if os.path.exists(os.path.join(c, 'interlocutors.json')):\n        data_root = c\n        break\n\n# Fallback: recursive search.\nif data_root is None:\n    for base in ['/content', '/content/drive/MyDrive']:\n        if not os.path.exists(base):\n            continue\n        for m in glob.glob(os.path.join(base, '**', 'interlocutors.json'), recursive=True):\n            if 'dialogues' in os.listdir(os.path.dirname(m)):\n                data_root = os.path.dirname(m)\n                break\n        if data_root:\n            break\n\nassert data_root, (\n    'interlocutors.json not found. Upload `real_persona_chat/` (with interlocutors.json + '\n    'dialogues/*.json) to one of:\\n  - ' + '\\n  - '.join(CANDIDATES)\n)\nprint('data_root:', data_root)\n\n# --- interlocutors -----------------------------------------------------------\nwith open(os.path.join(data_root, 'interlocutors.json'), encoding='utf-8') as f:\n    raw = json.load(f)\n\nif isinstance(raw, dict):\n    interlocutor_ds = []\n    for sid, info in raw.items():\n        if not isinstance(info, dict):\n            continue\n        info.setdefault('interlocutor_id', sid)\n        interlocutor_ds.append(info)\nelse:\n    interlocutor_ds = list(raw)\n\n# --- dialogues ---------------------------------------------------------------\ndialog_dir = os.path.join(data_root, 'dialogues')\nfiles = sorted(glob.glob(os.path.join(dialog_dir, '*.json')))\nassert files, f'no dialogue JSONs found in {dialog_dir}'\n\ndialogue_ds = []\nfor fp in files:\n    with open(fp, encoding='utf-8') as f:\n        d = json.load(f)\n    utts = d.get('utterances', [])\n    if isinstance(utts, list) and utts and isinstance(utts[0], dict):\n        d['utterances'] = {\n            'text':            [u.get('text', '') for u in utts],\n            'interlocutor_id': [u.get('interlocutor_id') or u.get('speaker_id') or '' for u in utts],\n        }\n    dialogue_ds.append(d)\n\nprint('dialogues:', len(dialogue_ds))\nprint('speakers :', len(interlocutor_ds))\nprint('\\ndialogue[0] keys    :', list(dialogue_ds[0].keys()))\nprint('interlocutor[0] keys:', list(interlocutor_ds[0].keys()))\nprint('personality fields  :', list(interlocutor_ds[0]['personality'].keys())[:6], '...')"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4 — Big Five ルックアップ + モノローグ CSV に変換\n",
    "\n",
    "各 (対話, 話者) ペアごとに、その話者の発言だけ `[SEP]` で結合した文字列を作り、ラベルは話者の Big Five (1-7) を付ける。\n",
    "→ `train.py` の `gene_data_from_csv` が期待する形式と一致。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "def speaker_id_of(row):\n",
    "    return row.get('interlocutor_id') or row.get('speaker_id')\n",
    "\n",
    "speaker_big5 = {}\n",
    "for row in interlocutor_ds:\n",
    "    sid = speaker_id_of(row)\n",
    "    p = row['personality']\n",
    "    speaker_big5[sid] = {\n",
    "        'n': float(p['BigFive_Neuroticism']),\n",
    "        'e': float(p['BigFive_Extraversion']),\n",
    "        'o': float(p['BigFive_Openness']),\n",
    "        'a': float(p['BigFive_Agreeableness']),\n",
    "        'c': float(p['BigFive_Conscientiousness']),\n",
    "    }\n",
    "print('speakers with Big Five:', len(speaker_big5))\n",
    "print('sample:', list(speaker_big5.items())[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rows = []\n",
    "skipped = 0\n",
    "for d in dialogue_ds:\n",
    "    did = d['dialogue_id']\n",
    "    utts = d['utterances']\n",
    "    texts = utts['text']\n",
    "    speakers = utts['interlocutor_id']\n",
    "    for sid in d['interlocutors']:\n",
    "        if sid not in speaker_big5:\n",
    "            skipped += 1\n",
    "            continue\n",
    "        my_texts = [t for t, s in zip(texts, speakers) if s == sid and isinstance(t, str) and t.strip()]\n",
    "        if not my_texts:\n",
    "            skipped += 1\n",
    "            continue\n",
    "        # train.py expects [CLS]...[SEP]...[SEP]... and replaces tags with tokenizer special tokens.\n",
    "        dialogue_text = '[CLS]' + '[SEP]'.join(my_texts)\n",
    "        b = speaker_big5[sid]\n",
    "        rows.append({\n",
    "            'file': f'{did}.csv',\n",
    "            'speaker_id': sid,\n",
    "            'dialogue': dialogue_text,\n",
    "            'n': b['n'], 'e': b['e'], 'o': b['o'], 'a': b['a'], 'c': b['c'],\n",
    "        })\n",
    "\n",
    "df = pd.DataFrame(rows)\n",
    "print(f'rows: {len(df)}   skipped: {skipped}')\n",
    "print('unique speakers:', df['speaker_id'].nunique())\n",
    "print(df.head(2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5 — 話者単位で train/valid/test 分割\n",
    "\n",
    "重要: 同じ話者を train と test の両方に入れると性格当てがリークする。**話者IDで分割**する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "rng = np.random.RandomState(42)\n",
    "speakers = sorted(df['speaker_id'].unique())\n",
    "rng.shuffle(speakers)\n",
    "n = len(speakers)\n",
    "train_sp = set(speakers[: int(0.8 * n)])\n",
    "valid_sp = set(speakers[int(0.8 * n): int(0.9 * n)])\n",
    "test_sp  = set(speakers[int(0.9 * n):])\n",
    "\n",
    "train_df = df[df['speaker_id'].isin(train_sp)].reset_index(drop=True)\n",
    "valid_df = df[df['speaker_id'].isin(valid_sp)].reset_index(drop=True)\n",
    "test_df  = df[df['speaker_id'].isin(test_sp)].reset_index(drop=True)\n",
    "\n",
    "print(f'train {len(train_df):>6} rows / {len(train_sp):>3} speakers')\n",
    "print(f'valid {len(valid_df):>6} rows / {len(valid_sp):>3} speakers')\n",
    "print(f'test  {len(test_df):>6} rows / {len(test_sp):>3} speakers')\n",
    "\n",
    "DATA_DIR = '/content/repo/data/rpc_monologue'\n",
    "os.makedirs(DATA_DIR, exist_ok=True)\n",
    "train_df.to_csv(f'{DATA_DIR}/train.csv', index=False)\n",
    "valid_df.to_csv(f'{DATA_DIR}/valid.csv', index=False)\n",
    "test_df.to_csv(f'{DATA_DIR}/test.csv',  index=False)\n",
    "print('CSVs written to', DATA_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6 — 学習\n",
    "\n",
    "**T4 メモ (16GB)**: LUKE-base + max_length=512 だと batch_size=128 (元設定) は OOM する。`batch_size=8` 程度から始める。\n",
    "\n",
    "**所要時間目安** (T4):\n",
    "- 1 epoch ≒ 10〜20分\n",
    "- 10 epoch で 2〜3時間\n",
    "- patience=3 で早期停止\n",
    "\n",
    "**A100/Pro なら** `batch_size=32`, max_epoch=15 が現実的。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir('/content/repo')\n",
    "\n",
    "LOG_FILE = f'{LOG_DIR}/v1.log'\n",
    "\n",
    "!python train.py \\\n",
    "  --data_folder ./data/rpc_monologue \\\n",
    "  --model_folder {MODEL_DIR} \\\n",
    "  --log_file {LOG_FILE} \\\n",
    "  --train_filename train.csv \\\n",
    "  --valid_filename valid.csv \\\n",
    "  --test_filename test.csv \\\n",
    "  --base_model_name studio-ousia/luke-japanese-base \\\n",
    "  --max_length 512 \\\n",
    "  --batch_size 8 \\\n",
    "  --lr 1e-5 \\\n",
    "  --warmup_steps 150 \\\n",
    "  --max_epoch 10 \\\n",
    "  --patience 3 \\\n",
    "  --critertion_type mae \\\n",
    "  --multitask 11111 \\\n",
    "  --multilinear 1 \\\n",
    "  --hidden_size 16 \\\n",
    "  --use_dropout 0 \\\n",
    "  --context 0 \\\n",
    "  --train_flag 1 \\\n",
    "  --test_flag 1 \\\n",
    "  --real_time_flag 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7 — DataParallel プレフィックス除去\n",
    "\n",
    "`train.py` は `nn.DataParallel(model)` で包んでから保存しているので、state_dict のキーが `module.base_model....` になっている。`predict_from_txt.py` は素の `MultiTaskModel` で読むので、`module.` プレフィックスを剥がして保存し直す。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "ckpt_path = f'{MODEL_DIR}/model.pt'\n",
    "state = torch.load(ckpt_path, map_location='cpu')\n",
    "if any(k.startswith('module.') for k in state.keys()):\n",
    "    state = {(k[len('module.'):] if k.startswith('module.') else k): v for k, v in state.items()}\n",
    "    torch.save(state, ckpt_path)\n",
    "    print('stripped module. prefix and re-saved')\n",
    "else:\n",
    "    print('no module. prefix to strip')\n",
    "\n",
    "size_mb = os.path.getsize(ckpt_path) / 1e6\n",
    "print(f'checkpoint: {ckpt_path}   ({size_mb:.1f} MB)')\n",
    "print('top 5 keys:', list(state.keys())[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 8 — 推論サニティチェック\n",
    "\n",
    "test split から 3 サンプル取って推定値と正解 (1-7) を比べる。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, torch.nn as nn\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "\n",
    "class MultiTaskModel(nn.Module):\n",
    "    def __init__(self, base, hidden_size=16):\n",
    "        super().__init__()\n",
    "        self.base_model = base\n",
    "        self.task_heads = nn.ModuleList([\n",
    "            nn.Sequential(\n",
    "                nn.Linear(base.config.hidden_size, hidden_size),\n",
    "                nn.ReLU(),\n",
    "                nn.Linear(hidden_size, 1),\n",
    "            ) for _ in range(5)\n",
    "        ])\n",
    "    def forward(self, input_ids, attention_mask):\n",
    "        pooled = self.base_model(input_ids=input_ids, attention_mask=attention_mask).pooler_output\n",
    "        return torch.cat([h(pooled) for h in self.task_heads], dim=1)\n",
    "\n",
    "tok = AutoTokenizer.from_pretrained('studio-ousia/luke-japanese-base')\n",
    "tok.add_tokens(['[SPK1]', '[SPK2]'])\n",
    "base = AutoModel.from_pretrained('studio-ousia/luke-japanese-base')\n",
    "base.resize_token_embeddings(len(tok))\n",
    "\n",
    "model = MultiTaskModel(base, hidden_size=16)\n",
    "model.load_state_dict(torch.load(ckpt_path, map_location='cpu'))\n",
    "model.eval()\n",
    "\n",
    "TRAITS = ['N', 'E', 'O', 'A', 'C']\n",
    "for i in range(min(3, len(test_df))):\n",
    "    s = test_df.iloc[i]\n",
    "    text = s['dialogue'].replace('[CLS]', tok.cls_token).replace('[SEP]', tok.sep_token)\n",
    "    enc = tok(text, max_length=512, padding='max_length', truncation=True,\n",
    "              return_tensors='pt', add_special_tokens=False)\n",
    "    with torch.no_grad():\n",
    "        pred = model(enc['input_ids'], enc['attention_mask'])[0].clamp(0, 1)\n",
    "    pred_1_7 = (pred * 6 + 1).tolist()\n",
    "    gold = [s['n'], s['e'], s['o'], s['a'], s['c']]\n",
    "    print(f'\\n--- speaker {s[\"speaker_id\"]} ---')\n",
    "    print('pred (1-7):', {t: round(v, 2) for t, v in zip(TRAITS, pred_1_7)})\n",
    "    print('gold (1-7):', {t: round(v, 2) for t, v in zip(TRAITS, gold)})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 完了\n",
    "\n",
    "- 学習済み `model.pt` は `/content/drive/MyDrive/Text2Persona/model/v1/model.pt` に保存済み\n",
    "- ローカル端末では `Personality-Recognition-on-RealPersonaChat/model/monologue_split_500k/model.pt` (空ファイル) を上書きすれば、既存の `predict_from_txt.py` でそのまま LINE `.txt` の推論ができる\n",
    "\n",
    "**精度の目安** (test split で見るべき指標)\n",
    "- N/E/O/A/C 各次元の Pearson r > 0.2 で『形になった』\n",
    "- > 0.3〜0.4 で論文水準\n",
    "- 過学習している場合は `--use_dropout 1`, `--max_epoch` を減らす, データ増強(500k化) を検討\n",
    "\n",
    "**次のステップ案**\n",
    "- 500-char 窓スライドで data augmentation して `monologue_split_500k` 相当を再現\n",
    "- 対話 (context=1) モードに切り替えて HC-GNN を試す (要 `torch_geometric`)\n",
    "- LINE `.txt` を実際に投入して推定値が直感と合うかを確認"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "machine_shape": "hm",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}