Donutモデルとは？

import json
import matplotlib.pyplot as plt
from datasets import load_dataset

# 1. データセットのロード
dataset = load_dataset("naver-clova-ix/cord-v2")

# 2. 訓練セットから最初のサンプル（インデックス0）を取得
sample_idx = 0
sample = dataset["train"][sample_idx]

image = sample["image"]
ground_truth_str = sample["ground_truth"]

# 3. JSON文字列を見やすく整形 (Pretty Print)
ground_truth_json = json.loads(ground_truth_str)
formatted_json = json.dumps(ground_truth_json, indent=4, ensure_ascii=False)

# 4. 画像の可視化
fig, ax = plt.subplots(figsize=(8, 10))
ax.imshow(image)
ax.axis("off") # 軸の目盛りを消す
ax.set_title(f"CORD-v2 Train Sample #{sample_idx}", fontsize=14)
plt.show()

# 5. 正解データ (Ground Truth) の表示
print("=== 正解データ (Ground Truth) ===")
print(formatted_json)

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms
from transformers import XLMRobertaTokenizer

## nn.TransformerEncoderによるシンプルなDonut Encoderの実装
class SimpleDonutEncoder(nn.Module):
    def __init__(self, hidden_dim=128, num_layers=2, nhead=4, max_seq_len=1024):
        super().__init__()
        self.patch_embed = nn.Conv2d(in_channels=3, out_channels=hidden_dim, kernel_size=16, stride=16)
        self.pos_embed = nn.Parameter(torch.randn(1, max_seq_len, hidden_dim))
        encoder_layer = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=nhead, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

    def forward(self, x):
        x = self.patch_embed(x) 
        x = x.flatten(2).transpose(1, 2) 
        seq_len = x.size(1)
        x = x + self.pos_embed[:, :seq_len, :]
        return self.transformer_encoder(x)

## nn.TransformerDecoderによるシンプルなDonut Decoderの実装
class SimpleDonutDecoder(nn.Module):
    def __init__(self, vocab_size, hidden_dim=128, max_seq_len=1024):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_dim)
        # テキスト側にも位置情報を追加
        self.pos_embed = nn.Parameter(torch.randn(1, max_seq_len, hidden_dim))

        self.decoder_layer = nn.TransformerDecoderLayer(d_model=hidden_dim, nhead=4, batch_first=True)
        self.transformer_decoder = nn.TransformerDecoder(self.decoder_layer, num_layers=2)

        self.lm_head = nn.Linear(hidden_dim, vocab_size)

    def forward(self, input_ids, encoder_hidden_states):
        seq_len = input_ids.size(1)
        device = input_ids.device

        tgt_embed = self.embedding(input_ids)
        tgt_embed = tgt_embed + self.pos_embed[:, :seq_len, :]

        # 将来トークンのカンニング防止マスク (Look-ahead mask / Causal mask)
        tgt_mask = torch.full((seq_len, seq_len), float('-inf'), device=device)
        tgt_mask = torch.triu(tgt_mask, diagonal=1)

        decoder_outputs = self.transformer_decoder(
            tgt=tgt_embed,
            memory=encoder_hidden_states,
            tgt_mask=tgt_mask
        )
        return self.lm_head(decoder_outputs)

class ConceptualDonutModel(nn.Module):
    def __init__(self, vocab_size, pad_token_id):
        super().__init__()
        self.encoder = SimpleDonutEncoder()
        self.decoder = SimpleDonutDecoder(vocab_size)
        # Padトークンを計算から除外する
        self.loss_fn = nn.CrossEntropyLoss(ignore_index=pad_token_id)

    def forward(self, image_tensors, decoder_input_ids, decoder_labels):
        encoder_outputs = self.encoder(image_tensors)
        logits = self.decoder(input_ids=decoder_input_ids, encoder_hidden_states=encoder_outputs)
        loss = self.loss_fn(logits.view(-1, logits.size(-1)), decoder_labels.view(-1))
        return loss

    def generate(self, image_tensors, start_prompt_id, eos_id, max_length=100):
        encoder_outputs = self.encoder(image_tensors)
        generated_ids = [start_prompt_id]

        for _ in range(max_length):
            input_ids = torch.tensor([generated_ids]).to(image_tensors.device)
            logits = self.decoder(input_ids, encoder_outputs)
            next_token_id = torch.argmax(logits[0, -1, :]).item()
            generated_ids.append(next_token_id)
            if next_token_id == eos_id:
                break
        return generated_ids

# Tokenizerの読み込み
tokenizer = XLMRobertaTokenizer.from_pretrained("naver-clova-ix/donut-base")
vocab_size = len(tokenizer)

# 今回は概念の確認のため、構造化データを400文字で切り詰める。
max_test_length = 400
short_ground_truth_str = ground_truth_str[:max_test_length]

# 構造化データ文字列のトークンIDへの変換
tokens = tokenizer.encode(short_ground_truth_str)

# GPUが利用できる環境ではcuda、CPUのみの場合はcpuを設定
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# モデルの初期化 (pad_token_id=1 を渡す)
model = ConceptualDonutModel(vocab_size=vocab_size, pad_token_id=tokenizer.pad_token_id)
model.to(device)

# optimizerの定義
optimizer = optim.AdamW(model.parameters(), lr=1e-3)

# 画像の前処理(256×256へのリサイズ、正規化の実施)
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
# 今回は訓練サンプルの0番目のデータに過学習するように学習させるが、モデルはバッチ入力が必要なため、バッチ次元として1次元目をbatch_size=1として追加
image_tensor = transform(image).unsqueeze(0).to(device)
tokens_tensor = torch.tensor(tokens).unsqueeze(0).to(device)

# Teacher Forcingのためのラベルデータの追加(tokens_tensorの0番目にはbos_idが設定されている。)
decoder_input_ids = tokens_tensor[:, :-1]
decoder_labels = tokens_tensor[:, 1:]

# モデルの学習
model.train()
epochs = 300

print("モデルの学習開始")
for epoch in range(1, epochs + 1):
    optimizer.zero_grad()
    loss = model(image_tensor, decoder_input_ids, decoder_labels)
    loss.backward()
    optimizer.step()

    if epoch % 20 == 0 or epoch == 1:
        print(f"Epoch {epoch:03d}/{epochs} - Loss: {loss.item():.4f}")
print("学習完了")

print("モデル評価")
model.eval()

with torch.no_grad():
    predicted_ids = model.generate(
        image_tensors=image_tensor,
        start_prompt_id=tokenizer.bos_token_id, # bosトークンID
        eos_id=tokenizer.eos_token_id,          # eosトークンID
        max_length=len(tokens) + 10
    )

# 特殊トークンを除外してテキストに復元
predicted_text = tokenizer.decode(predicted_ids, skip_special_tokens=True)
# 正解側も、特殊トークン抜きの純粋なテキストにするため一度デコードして比較基準にします
target_text = tokenizer.decode(tokens, skip_special_tokens=True)

print("【モデルの予測結果】")
print(predicted_text)
print("-" * 30)
print("【本来の正解データ(400文字に切り詰め】")
print(target_text)

if predicted_text == target_text:
    print("予測と正解データが一致しました。")
else:
    print("予測と正解データが一致しませんでした。")

モデルの学習開始
Epoch 001/300 - Loss: 10.9835
Epoch 020/300 - Loss: 3.9394
Epoch 040/300 - Loss: 0.6503
Epoch 060/300 - Loss: 0.1062
Epoch 080/300 - Loss: 0.0493
Epoch 100/300 - Loss: 0.0337
Epoch 120/300 - Loss: 0.0261
Epoch 140/300 - Loss: 0.0214
Epoch 160/300 - Loss: 0.0179
Epoch 180/300 - Loss: 0.0153
Epoch 200/300 - Loss: 0.0131
Epoch 220/300 - Loss: 0.0115
Epoch 240/300 - Loss: 0.0100
Epoch 260/300 - Loss: 0.0088
Epoch 280/300 - Loss: 0.0080
Epoch 300/300 - Loss: 0.0073
学習完了
モデル評価
【モデルの予測結果】
{"gt_parse": {"menu": [{"nm": "Nasi Campur Bali", "cnt": " x", "price": "75,000"}, {"nm": "Bbk Bengil Nasi", "cnt": " x", "price": "125,000"}, {"nm": "MilkShake Starwb", "cnt": " x", "price": "37,000"}, {"nm": "Ice Lemon Tea", "cnt": " x", "price": "24,000"}, {"nm": "Nasi Ayam Dewata", "cnt": " x", "price": "70,000"}, {"nm": "Free Ice Tea", "cnt": "3 x", "price": "0"}, {"nm": "Organic Green S
------------------------------
【本来の正解データ(400文字に切り詰め】
{"gt_parse": {"menu": [{"nm": "Nasi Campur Bali", "cnt": " x", "price": "75,000"}, {"nm": "Bbk Bengil Nasi", "cnt": " x", "price": "125,000"}, {"nm": "MilkShake Starwb", "cnt": " x", "price": "37,000"}, {"nm": "Ice Lemon Tea", "cnt": " x", "price": "24,000"}, {"nm": "Nasi Ayam Dewata", "cnt": " x", "price": "70,000"}, {"nm": "Free Ice Tea", "cnt": "3 x", "price": "0"}, {"nm": "Organic Green S
予測と正解データが一致しました。

from transformers import DonutProcessor, VisionEncoderDecoderModel

# 1. デバイスの設定
device = "cuda" if torch.cuda.is_available() else "cpu"

# 2. モデルとプロセッサの読み込み (CORDデータセット用に学習済みのモデルを使用)
model_id = "naver-clova-ix/donut-base-finetuned-cord-v2"
processor = DonutProcessor.from_pretrained(model_id)
model = VisionEncoderDecoderModel.from_pretrained(model_id)
model.to(device)

# 3. データセットのロード（テストセットから最初の1件を取得）
test_sample = dataset["test"][0]
image = test_sample["image"]

# 4. タスクプロンプトの準備と前処理
# CORD用のファインチューニングモデルは "<s_cord-v2>" というプロンプトから生成を開始します
task_prompt = "<s_cord-v2>"
decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
pixel_values = processor(image, return_tensors="pt").pixel_values

# 5. モデルによる生成（推論）の実行
print("モデルによる推論を実行中...")
outputs = model.generate(
    pixel_values.to(device),
    decoder_input_ids=decoder_input_ids.to(device),
    max_length=model.decoder.config.max_position_embeddings,
    pad_token_id=processor.tokenizer.pad_token_id,
    eos_token_id=processor.tokenizer.eos_token_id,
    use_cache=True,
    bad_words_ids=[[processor.tokenizer.unk_token_id]],
    return_dict_in_generate=True,
)

# 6. 出力トークンのデコードとJSONへのパース
sequence = processor.batch_decode(outputs.sequences)[0]
# 不要な特殊トークン（EOSやPAD）を削除
sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
# タスクプロンプト以降の文字列を抽出
sequence = sequence.split("<s_cord-v2>", 1)[-1].strip()

# DonutProcessorの便利機能：出力されたXML風のタグ構造をPythonの辞書（JSON）に自動変換
parsed_json = processor.token2json(sequence)

print("=== 実際のDonutモデルによる推論結果 ===")
print(json.dumps(parsed_json, indent=4, ensure_ascii=False))

# 7. 画像の可視化
fig, ax = plt.subplots(figsize=(8, 10))
ax.imshow(image)
ax.axis("off") # 軸の目盛りを消す
ax.set_title(f"CORD-v2 Test Sample #0", fontsize=14)
plt.show()

モデルによる推論を実行中...
=== 実際のDonutモデルによる推論結果 ===
{
    "menu": {
        "nm": "- TICKET CP",
        "num": "901016",
        "unitprice": "60.000",
        "cnt": "2",
        "price": "60.000"
    },
    "sub_total": {
        "subtotal_price": "-60.000",
        "tax_price": "5,455"
    },
    "total": {
        "total_price": "60.000",
        "emoneyprice": "60.000",
        "menuqty_cnt": "2.00"
    }
}