Nougatモデルとは？

from datasets import load_dataset
import matplotlib.pyplot as plt

# Hugging Faceから 'yuntian-deng/im2latex-100k'をロード
dataset = load_dataset("yuntian-deng/im2latex-100k")


# 訓練セットの最初のサンプルを取得
sample_idx = 0

sample = dataset["train"][sample_idx]
raw_image = sample["image"]      # PIL Image オブジェクト
latex_formula = sample["formula"] # 正解のLaTeX文字列

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

print(f"正解データ:LaTeX数式")
print(latex_formula)

正解データ:LaTeX数式
\widetilde \gamma _ { \mathrm { h o p f } } \simeq \sum _ { n > 0 } \widetilde { G } _ { n } { \frac { ( - a ) ^ { n } } { 2 ^ { 2 n - 1 } } }

import torch
import torch.nn as nn


class SimpleNougatEncoder(nn.Module):
    def __init__(self, hidden_dim=128, num_layers=2, nhead=4, max_seq_len=1024):
        super().__init__()
        # 数式画像は横長（例: 幅320px、高さ64px）が多いため、パッチサイズ16x16で分割
        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によるシンプルなNougat Decoderの実装
class SimpleNougatDecoder(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, :]

        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 ConceptualNougatModel(nn.Module):
    def __init__(self, vocab_size, pad_token_id):
        super().__init__()
        self.encoder = SimpleNougatEncoder()
        self.decoder = SimpleNougatDecoder(vocab_size)
        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=150):
        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

import torch.optim as optim
from torchvision import transforms
from transformers import NougatProcessor


# 本物のNougatトークナイザーをロード
processor = NougatProcessor.from_pretrained("facebook/nougat-base")
tokenizer = processor.tokenizer
vocab_size = len(tokenizer)

# LaTeX文字列のトークンID化
tokens = tokenizer.encode(latex_formula)

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

model = ConceptualNougatModel(vocab_size=vocab_size, pad_token_id=tokenizer.pad_token_id)
model.to(device)
optimizer = optim.AdamW(model.parameters(), lr=5e-4)

# 数式画像用の前処理（横長のアスペクト比にリサイズしてテンソル化）
transform = transforms.Compose([
    transforms.Resize((64, 320)), # IM2LATEX標準の数式サイズ感に合わせる
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

# PIL画像をRGBに変換してから前処理を適用（バッチ次元を追加）
image_tensor = transform(raw_image.convert("RGB")).unsqueeze(0).to(device)
tokens_tensor = torch.tensor(tokens).unsqueeze(0).to(device)

# Teacher Forcing用の入出力ペア
decoder_input_ids = tokens_tensor[:, :-1]
decoder_labels = tokens_tensor[:, 1:]


model.train()
epochs = 200
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,
        eos_id=tokenizer.eos_token_id,
        max_length=len(tokens) + 10
    )

# 特殊トークンを除外してテキストに復元
predicted_latex = tokenizer.decode(predicted_ids, skip_special_tokens=True)
target_latex = tokenizer.decode(tokens, skip_special_tokens=True)

print("【モデルの予測結果】")
print(target_latex)
print("【本来の正解データ")
print(predicted_latex)

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

モデルの学習開始
Epoch 001/200 - Loss: 10.8126
Epoch 020/200 - Loss: 5.6297
Epoch 040/200 - Loss: 2.7499
Epoch 060/200 - Loss: 1.0036
Epoch 080/200 - Loss: 0.3588
Epoch 100/200 - Loss: 0.1623
Epoch 120/200 - Loss: 0.1013
Epoch 140/200 - Loss: 0.0728
Epoch 160/200 - Loss: 0.0572
Epoch 180/200 - Loss: 0.0474
Epoch 200/200 - Loss: 0.0392
学習完了
モデル評価
【モデルの予測結果】
\widetilde \gamma _ { \mathrm { h o p f } } \simeq \sum _ { n > 0 } \widetilde { G } _ { n } { \frac { ( - a ) ^ { n } } { 2 ^ { 2 n - 1 } } }
【本来の正解データ
\widetilde \gamma _ { \mathrm { h o p f } } \simeq \sum _ { n > 0 } \widetilde { G } _ { n } { \frac { ( - a ) ^ { n } } { 2 ^ { 2 n - 1 } } }
予測と正解データが一致しました。

import pypdfium2 as pdfium
import torch
from PIL import Image
from transformers import VisionEncoderDecoderModel, NougatTokenizerFast
from torchvision import transforms


device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "facebook/nougat-base"


print("モデルとトークナイザーの読み込み")
tokenizer = NougatTokenizerFast.from_pretrained(model_id)
model = VisionEncoderDecoderModel.from_pretrained(model_id).to(device)
model.eval()


# PDFをPIL画像に変換する関数(scale=2に設定することで、数式や細かい文字をクリアにレンダリング)
def convert_pdf_page_to_image(pdf_path, page_idx=0, scale=2):
    pdf = pdfium.PdfDocument(pdf_path)
    if page_idx >= len(pdf):
        raise ValueError(f"指定されたページ番号 {page_idx} はPDFの総ページ数（{len(pdf)}ページ）を超えています。")

    page = pdf[page_idx]
    bitmap = page.render(scale=scale)
    pil_img = bitmap.to_pil()
    return pil_img.convert("RGB")


nougat_transform = transforms.Compose([
    transforms.Resize((896, 672)),  # Nougatの規定入力サイズ (height, width)
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406], # 標準的な輝度平均
        std=[0.229, 0.224, 0.225]   # 標準偏差
    )
])

# ※検証に用いるPMCのpdfファイル
# 論文名: Machine learning in medicine: a practical introduction
# 著者: Jenni A. M. Sidey-Gibbons and Chris J. Sidey-Gibbons
# 出典: BMC Medical Research Methodology (2019) 19:64 (https://doi.org/10.1186/s12874-019-0681-4)
# ライセンス: 本論文は [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/deed.ja) ライセンスの下で公開されており、原著作者のクレジットを表示することで複製・再配布が許可されています。
pdf_filename = "12874_2019_Article_681.pdf"

# 論文の「1ページ目（インデックス0）」を画像化
image = convert_pdf_page_to_image(pdf_filename, page_idx=0, scale=2)

# テンソル化、バッチ次元の追加
pixel_values = nougat_transform(image).unsqueeze(0).to(device)

print(f"『{pdf_filename}』を解析...")

with torch.no_grad():
    outputs = model.generate(
        pixel_values,
        min_length=1,
        max_new_tokens=2000, # 1ページ全体の文章と数式を出力するため長めに設定
        decoder_start_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
    )

# トークン列をMarkdownテキストに直接デコード
predicted_markdown = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
print("解析結果")
print(predicted_markdown)