import math
import re
import torch
from transformers import GPT2LMHeadModel, GPT2TokenizerFast

# モデルの準備
MODEL_NAME = "gpt2"
tokenizer = GPT2TokenizerFast.from_pretrained(MODEL_NAME)
model = GPT2LMHeadModel.from_pretrained(MODEL_NAME)
model.eval()

# 文を受け取り、単語ごとのサプライザル (単位: bit) を辞書形式で返す関数
def get_word_surprisals(sentence: str) -> dict:
    
    inputs = tokenizer(sentence, return_tensors="pt")
    input_ids = inputs["input_ids"]

    # モデルに文を渡し、各トークン位置での全語彙に対する予測スコア（logits）を取得
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits

    # 「t番目までの文脈」から「t+1番目の単語」を予測するため、1つずらす（Shift）
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = input_ids[:, 1:].contiguous()

    # Log-Softmaxで対数確率 ln(P) を求め、実際に出現した単語の確率を抽出
    log_probs = torch.nn.functional.log_softmax(shift_logits, dim=-1)
    target_log_probs = log_probs.gather(2, shift_labels.unsqueeze(-1)).squeeze(-1)[0]

    # 自然対数 ln(P) を底2の対数 -log2(P) に変換して bits 単位のサプライザルにする
    token_surprisals = [-lp.item() / math.log(2) for lp in target_log_probs]

    # トークン（サブワード）を元の単語単位にまとめる
    encodings = tokenizer(sentence, return_offsets_mapping=True)
    offset_mapping = encodings["offset_mapping"]

    word_matches = list(re.finditer(r"\b[\w'-]+\b", sentence))
    result = {}

    for match in word_matches:
        w_start, w_end = match.span()
        word = match.group()
        total_bits = 0.0
        matched = False

        # offset_mapping[0] は文頭（先行文脈なし）のため、2番目のトークン以降を照合
        for i, (t_start, t_end) in enumerate(offset_mapping[1:], start=0):
            if max(w_start, t_start) < min(w_end, t_end):
                total_bits += token_surprisals[i]
                matched = True

        result[word] = round(total_bits, 2) if matched else None

    return result

sentence = "There are two man at the door."
surprisals = get_word_surprisals(sentence)

print(f"{'Word':<10} | {'Surprisal (bits)':<18} | {'Graph'}")
print("-" * 45)
for word, val in surprisals.items():
    if val is not None:
        bar = "■" * int(val)
        print(f"{word:<10} | {val:<18.2f} | {bar}")
    else:
        print(f"{word:<10} | {'(Initial word)':<18} | ")

Word       | Surprisal (bits)   | Graph
---------------------------------------------
There      | (Initial word)     | 
are        | 4.97               | ■■■■
two        | 3.95               | ■■■
man        | 14.90              | ■■■■■■■■■■■■■■
at         | 8.72               | ■■■■■■■■
the        | 1.23               | ■
door       | 5.55               | ■■■■■

# 実験：文法的に不自然な箇所を持つ文で試す
sentence = "She goes there yesterday"
surprisals = get_word_surprisals(sentence)

# 結果を見やすく整形して出力
print(f"{'Word':<10} | {'Surprisal (bits)':<18} | {'Graph'}")
print("-" * 45)
for word, val in surprisals.items():
    if val is not None:
        bar = "■" * int(val)
        print(f"{word:<10} | {val:<18.2f} | {bar}")
    else:
        print(f"{word:<10} | {'(Initial word)':<18} | ")

Word       | Surprisal (bits)   | Graph
---------------------------------------------
She        | (Initial word)     | 
goes       | 10.68              | ■■■■■■■■■■
there      | 8.60               | ■■■■■■■■
yesterday  | 12.81              | ■■■■■■■■■■■■

学習者言語の分析（応用）1（第6回）¶

6.1 導入¶

6.2 サプライザルに関する背景¶

6.2.1 情報理論¶

6.2.2 Hale (2001)のA Probabilistic Earley Parser as a Psycholinguistic Model¶

6.3 サプライザルの計算¶

6.4 計算してみよう¶