Fine-Tuning Gemma 4 31B on CORD-v2 Receipts — End-to-End Guide

1
import json
2
from datasets import load_dataset
3
4
dataset = load_dataset("naver-clova-ix/cord-v2")
5
6
print(f"Train      : {len(dataset['train'])} samples")
7
print(f"Validation : {len(dataset['validation'])} samples")
8
print(f"Test       : {len(dataset['test'])} samples")
9
10
# Inspect a sample
11
ex = dataset["train"][0]
12
gt = json.loads(ex["ground_truth"])["gt_parse"]
13
print(f"\nGround-truth keys: {list(gt.keys())}")
14
15
for key in gt:
16
    if isinstance(gt[key], list):
17
        print(f"\n  {key} ({len(gt[key])} items):")
18
        for item in gt[key][:2]:
19
            print(f"    {item}")
20
    else:
21
        print(f"\n  {key}: {gt[key]}")

1
from nemo_automodel.components.datasets.vlm.utils import json2token
2
3
token_seq = json2token(gt, sort_json_key=True)
4
print(f"Token sequence (first 300 chars):\n  {token_seq[:300]}...")
5
print(f"\nTotal length: {len(token_seq)} chars")

1
import os
2
import json
3
import torch
4
from transformers import AutoProcessor
5
from nemo_automodel import NeMoAutoModelForImageTextToText
6
from nemo_automodel.components.datasets.vlm.utils import json2token
7
from datasets import load_dataset
8
9
# --- Helpers ---
10
11
def compute_ned(pred: str, target: str) -> float:
12
    """Normalized Edit Distance (0 = perfect match, 1 = completely different)."""
13
    m, n = len(pred), len(target)
14
    if max(m, n) == 0:
15
        return 0.0
16
    dp = list(range(n + 1))
17
    for i in range(1, m + 1):
18
        prev, dp[0] = dp[0], i
19
        for j in range(1, n + 1):
20
            tmp = dp[j]
21
            dp[j] = prev if pred[i - 1] == target[j - 1] else 1 + min(dp[j], dp[j - 1], prev)
22
            prev = tmp
23
    return dp[n] / max(m, n)
24
25
def run_gemma4_inference(model, processor, pil_image, prompt="Describe this image.",
26
                         max_new_tokens=1024):
27
    """Run Gemma 4 inference on a single image."""
28
    messages = [
29
        {
30
            "role": "user",
31
            "content": [
32
                {"type": "image", "image": pil_image},
33
                {"type": "text", "text": prompt},
34
            ],
35
        },
36
    ]
37
    inputs = processor.apply_chat_template(
38
        messages,
39
        tokenize=True,
40
        add_generation_prompt=True,
41
        return_tensors="pt",
42
        return_dict=True,
43
    ).to(model.device)
44
45
    with torch.inference_mode():
46
        outputs = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
47
48
    generated_text = processor.decode(outputs[0], skip_special_tokens=True)
49
    prompt_length = len(processor.decode(inputs["input_ids"][0], skip_special_tokens=True))
50
    return generated_text[prompt_length:].strip()
51
52
def evaluate_receipts(model, processor, test_dataset, n_samples=20):
53
    """Evaluate model on receipt test set, return avg NED and per-sample results."""
54
    model.eval()
55
    results = []
56
    n = min(n_samples, len(test_dataset))
57
    for i in range(n):
58
        ex = test_dataset[i]
59
        gt = json.loads(ex["ground_truth"])["gt_parse"]
60
        target = json2token(gt, sort_json_key=True)
61
        pred = run_gemma4_inference(model, processor, ex["image"])
62
        ned = compute_ned(pred, target)
63
        results.append({"idx": i, "ned": ned, "pred": pred, "target": target, "gt": gt})
64
        print(f"  Sample {i:2d}: NED = {ned:.4f}")
65
    avg_ned = sum(r["ned"] for r in results) / len(results)
66
    print(f"\n  Average NED: {avg_ned:.4f}")
67
    return avg_ned, results
68
69
# --- Load base model ---
70
71
MODEL_PATH = "google/gemma-4-31B-it"
72
73
processor = AutoProcessor.from_pretrained(MODEL_PATH)
74
base_model = NeMoAutoModelForImageTextToText.from_pretrained(
75
    MODEL_PATH,
76
    torch_dtype=torch.bfloat16,
77
    use_liger_kernel=True,
78
    attn_implementation="flash_attention_2",
79
    text_config={"use_cache": False},
80
).eval().to("cuda")
81
82
print(f"Parameters: {sum(p.numel() for p in base_model.parameters()):,}")
83
84
# --- Evaluate ---
85
86
dataset = load_dataset("naver-clova-ix/cord-v2")
87
print("\nEvaluating base model on receipt test set:")
88
base_avg_ned, base_results = evaluate_receipts(base_model, processor, dataset["test"])

1
import json
2
import os
3
import torch
4
from datasets import load_dataset
5
from transformers import AutoProcessor, AutoModelForMultimodalLM
6
from nemo_automodel.components.datasets.vlm.utils import json2token
7
8
# Paths
9
BASE_MODEL = "google/gemma-4-31B-it"
10
CKPT_DIR = "vlm_checkpoints/gemma4_31b_cord_v2"
11
best_ckpt = os.path.realpath(os.path.join(CKPT_DIR, "LOWEST_VAL"))
12
consolidated = os.path.join(best_ckpt, "model", "consolidated")
13
14
# Load processor from base model, model from fine-tuned checkpoint
15
processor = AutoProcessor.from_pretrained(BASE_MODEL)
16
model = AutoModelForMultimodalLM.from_pretrained(
17
    consolidated,
18
    dtype=torch.bfloat16,
19
    device_map="auto",
20
).eval()
21
22
# Evaluate on test set
23
dataset = load_dataset("naver-clova-ix/cord-v2")
24
print("Evaluating fine-tuned model:")
25
ft_avg_ned, ft_results = evaluate_receipts(model, processor, dataset["test"])

1
import re
2
3
def token2json(token_seq):
4
    """Convert a token sequence back to a JSON-like dict."""
5
    result = {}
6
    pattern = r"<s_(\w+)>(.*?)</s_\1>"
7
    matches = re.findall(pattern, token_seq, re.DOTALL)
8
    for key, value in matches:
9
        if "<sep/>" in value:
10
            items = value.split("<sep/>")
11
            result[key] = [token2json(item) if "<s_" in item else item for item in items]
12
        elif "<s_" in value:
13
            result[key] = token2json(value)
14
        else:
15
            result[key] = value
16
    return result
17
18
parsed = token2json(prediction)
19
print(json.dumps(parsed, indent=2))