data/Models/Text2History/TinySwallow.py

"""TinySwallow で LINE 会話履歴を要約する。"""

from __future__ import annotations

import argparse
import re
from pathlib import Path
from typing import Any

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


MODEL_NAME = "SakanaAI/TinySwallow-1.5B-Instruct"
__all__ = ["summarize_history_from_txt"]

DEFAULT_HISTORY_PATH = (
    Path(__file__).resolve().parents[2]
    / "text_save"
    / "[LINE]Sora Kura ٩( ᐛ )و.txt"
)

SYSTEM_PROMPT = (
    "あなたは会話履歴を1〜2文で要約するアシスタントです。"
    "直近の発言の文字列をそのままコピーせず、自分の言葉で言い換えて要約してください。"
    "過去要約に書かれた事実は引き継ぎ、最新発言で新たに分かった事実や話題の変化を反映してください。"
    "最新の発言が過去要約の話題と無関係な新しい話題に切り替わった場合は、"
    "古い話題は省略し、新しい話題を中心とした要約に切り替えてください。"
    "事実を改変したり、書かれていない情報を補ったりしてはいけません。"
    "出力は要約本文のみとし、'user:' 'assistant:' のようなロール名や前置き・箇条書きは書かないでください。"
)


def _validate_positive_int(name: str, value: int | None) -> None:
    if value is not None and value <= 0:
        raise ValueError(f"{name} must be a positive integer")


def pick_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def pick_dtype(device: str) -> torch.dtype:
    if device.startswith("cuda") or device == "mps":
        return torch.float16
    return torch.float32


LINE_MESSAGE_RE = re.compile(r"^(\d{2}:\d{2})\s+(\S+)\s+(.+)$")


def parse_line_export(
    path: Path,
    self_name: str,
    *,
    max_lines: int | None = None,
    use_tail: bool = True,
) -> list[dict[str, str]]:
    """LINE の .txt からメッセージ行を取り出す。

    max_lines は .txt の生の行数に対する上限。None の場合は全行を読む。
    use_tail=True のときは末尾 max_lines 行、False のときは先頭 max_lines 行を使う。
    """
    _validate_positive_int("max_lines", max_lines)

    lines = path.read_text(encoding="utf-8").splitlines()
    if max_lines is not None:
        lines = lines[-max_lines:] if use_tail else lines[:max_lines]

    history: list[dict[str, str]] = []
    for raw in lines:
        match = LINE_MESSAGE_RE.match(raw.strip())
        if not match:
            continue
        _, sender, text = match.groups()
        role = "user" if sender == self_name else "assistant"
        history.append({"role": role, "text": text})
    return history


def bootstrap_summary(
    history: list[dict[str, str]],
    batch_size: int,
    tokenizer,
    model,
    device: str,
    *,
    verbose: bool = False,
) -> str:
    _validate_positive_int("batch_size", batch_size)

    summary = ""
    total = len(history)
    for i in range(0, total, batch_size):
        chunk = history[i : i + batch_size]
        summary = summarize(summary, chunk, tokenizer, model, device)
        if verbose:
            print(f"  [bootstrap {min(i + batch_size, total)}/{total}] {summary}")
    return summary


def summarize(old_summary: str, recent_messages, tokenizer, model, device: str) -> str:
    recent_text = "\n".join(
        [f"{m['role']}: {m['text']}" for m in recent_messages]
    )

    user_content = (
        f"過去要約:\n{old_summary if old_summary else '(なし)'}\n\n"
        f"直近会話:\n{recent_text}\n\n"
        "上記を踏まえ、現在の会話の流れを1文で要約してください。"
    )

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_content},
    ]

    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = tokenizer([text], return_tensors="pt").to(device)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=100,
            do_sample=False,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
        )

    generated = output_ids[0][inputs.input_ids.shape[-1]:]
    return tokenizer.decode(generated, skip_special_tokens=True).strip()


def summarize_history_from_txt(
    txt_path: str | Path,
    *,
    self_name: str = "Rayta",
    batch_size: int = 10,
    max_lines: int | None = None,
    use_tail: bool = True,
    tokenizer: Any | None = None,
    model: Any | None = None,
    device: str | None = None,
    verbose: bool = False,
) -> str:
    """LINE 履歴 .txt を要約し、要約文字列だけを返す。

    この関数は history_summary.json などのファイルを作成・更新しない。

    Args:
        txt_path: 入力する LINE 履歴 .txt のパス。
        self_name: 自分の表示名。この送信者を user、それ以外を assistant として扱う。
        batch_size: モデルに一度に渡して要約するメッセージ行数。
        max_lines: .txt から使う最大行数。None の場合は全行を使う。
        use_tail: max_lines 指定時、末尾から使うなら True、先頭から使うなら False。
        tokenizer: 事前ロード済み tokenizer。model とセットで渡す。
        model: 事前ロード済み model。tokenizer とセットで渡す。
        device: "cpu" / "mps" / "cuda" など。None の場合は自動選択。
        verbose: True の場合、バッチごとの途中要約を表示する。
    """
    _validate_positive_int("batch_size", batch_size)
    _validate_positive_int("max_lines", max_lines)

    if (tokenizer is None) != (model is None):
        raise ValueError("tokenizer and model must be provided together")

    selected_device = device or pick_device()
    if tokenizer is None or model is None:
        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            torch_dtype=pick_dtype(selected_device),
        ).to(selected_device)
        model.eval()

    history = parse_line_export(
        Path(txt_path),
        self_name,
        max_lines=max_lines,
        use_tail=use_tail,
    )
    if not history:
        return ""

    return bootstrap_summary(
        history,
        batch_size,
        tokenizer,
        model,
        selected_device,
        verbose=verbose,
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="LINE履歴を読み込んで要約文字列を出力する。"
    )
    parser.add_argument(
        "--history",
        type=Path,
        default=DEFAULT_HISTORY_PATH,
        help="LINE履歴(.txt)のパス。",
    )
    parser.add_argument(
        "--self-name",
        default="Rayta",
        help="自分の表示名。これに一致する送信者をuserとして扱う。",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=10,
        help="モデルに一度に渡して要約するメッセージ行数。",
    )
    parser.add_argument(
        "--max-lines",
        type=int,
        default=None,
        help="入力 .txt から使う最大行数。省略時は全行。",
    )
    parser.add_argument(
        "--use-head",
        action="store_true",
        help="--max-lines 指定時に末尾ではなく先頭から読む。",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    device = pick_device()
    print(f"Model:  {MODEL_NAME}")
    print(f"Device: {device}\n")

    if not args.history.exists():
        print(f"履歴ファイルが見つかりません: {args.history}\n")
        return

    print(f"履歴ファイル読み込み中: {args.history}")
    summary = summarize_history_from_txt(
        args.history,
        self_name=args.self_name,
        batch_size=args.batch_size,
        max_lines=args.max_lines,
        use_tail=not args.use_head,
        device=device,
        verbose=True,
    )
    print(f"\n要約: {summary}\n")


if __name__ == "__main__":
    main()