Fine-Tune 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
26
def run_gemma4_inference(model, processor, pil_image, prompt="Describe this image.",
27
                         max_new_tokens=1024):
28
    """Run Gemma 4 inference on a single image."""
29
    messages = [
30
        {
31
            "role": "user",
32
            "content": [
33
                {"type": "image", "image": pil_image},
34
                {"type": "text", "text": prompt},
35
            ],
36
        },
37
    ]
38
    inputs = processor.apply_chat_template(
39
        messages,
40
        tokenize=True,
41
        add_generation_prompt=True,
42
        return_tensors="pt",
43
        return_dict=True,
44
    ).to(model.device)
45
46
    with torch.inference_mode():
47
        outputs = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
48
49
    generated_text = processor.decode(outputs[0], skip_special_tokens=True)
50
    prompt_length = len(processor.decode(inputs["input_ids"][0], skip_special_tokens=True))
51
    return generated_text[prompt_length:].strip()
52
53
54
def evaluate_receipts(model, processor, test_dataset, n_samples=20):
55
    """Evaluate model on receipt test set, return avg NED and per-sample results."""
56
    model.eval()
57
    results = []
58
    n = min(n_samples, len(test_dataset))
59
    for i in range(n):
60
        ex = test_dataset[i]
61
        gt = json.loads(ex["ground_truth"])["gt_parse"]
62
        target = json2token(gt, sort_json_key=True)
63
        pred = run_gemma4_inference(model, processor, ex["image"])
64
        ned = compute_ned(pred, target)
65
        results.append({"idx": i, "ned": ned, "pred": pred, "target": target, "gt": gt})
66
        print(f"  Sample {i:2d}: NED = {ned:.4f}")
67
    avg_ned = sum(r["ned"] for r in results) / len(results)
68
    print(f"\n  Average NED: {avg_ned:.4f}")
69
    return avg_ned, results
70
71
# --- Load base model ---
72
73
MODEL_PATH = "google/gemma-4-31B-it"
74
75
processor = AutoProcessor.from_pretrained(MODEL_PATH)
76
base_model = NeMoAutoModelForImageTextToText.from_pretrained(
77
    MODEL_PATH,
78
    torch_dtype=torch.bfloat16,
79
    use_liger_kernel=True,
80
    attn_implementation="flash_attention_2",
81
    text_config={"use_cache": False},
82
).eval().to("cuda")
83
84
print(f"Parameters: {sum(p.numel() for p in base_model.parameters()):,}")
85
86
# --- Evaluate ---
87
88
dataset = load_dataset("naver-clova-ix/cord-v2")
89
print("\nEvaluating base model on receipt test set:")
90
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))