Fine-Tuning di T5Gemma-2

T5Gemma 2 è una famiglia di modelli encoder–decoder leggeri e open-weight di Google, basati su Gemma 3, che supportano input multilingue e multimodali. 

Con una finestra di contesto fino a 128K in oltre 140 lingue e scelte progettuali efficienti in termini di parametri come tied embeddings e merged attention, questi modelli sono ideali per il text generation e la comprensione delle immagini, rimanendo abbastanza piccoli da poter essere eseguiti su un laptop.

In questo tutorial impareremo a fare fine-tuning di un modello encoder–decoder su un dataset di OCR LaTeX. L’obiettivo è ottenere buone prestazioni usando un numero minimo di campioni di training. 

1. Configurare l’ambiente GPU A100 su RunPod

Sebbene sia possibile effettuare il fine-tuning di questo modello su Kaggle o Google Colab, spesso ciò porta a sessioni instabili, disconnessioni e training sensibilmente più lenti. Per evitare questi attriti e mantenere un setup semplice e affidabile, useremo una GPU NVIDIA A100.

Puoi noleggiare una A100 su RunPod a circa 1,39 $/ora, e l’intero processo di training descritto in questo tutorial dovrebbe richiedere ben meno di 30 minuti. Questo setup ti garantisce prestazioni costanti senza dover lottare con vincoli di memoria.

Inizia andando su RunPod e creando un nuovo pod usando l’ultima immagine PyTorch. Seleziona una macchina 1× A100.

Poi modifica la configurazione del pod e aggiungi una variabile d’ambiente chiamata HF_TOKEN. Questo token è necessario per:

• Caricare modelli protetti su Hugging Face
• Pubblicare il tuo modello fine-tunato su Hugging Face Hub

Quando il pod è pronto, avvia il notebook e installa i pacchetti Python richiesti. Assicurati di usare l’ultima versione di transformers.

!pip -q install -U accelerate datasets pillow sentencepiece safetensors
!pip install --quiet "transformers==5.0.0rc1"
!pip install --quiet --no-deps trlOfferto da

Ora importa le librerie e le utility che useremo nel notebook.

import torch
from datasets import load_dataset
from transformers import AutoProcessor, AutoModelForSeq2SeqLM
from transformers import Seq2SeqTrainingArguments, Seq2SeqTrainer, set_seedOfferto da

Prima del training applichiamo alcune configurazioni ottimizzate per GPU A100. Impostare un seed garantisce la riproducibilità e abilitare TF32 migliora le prestazioni senza influire sulla stabilità del training in bf16.

set_seed(42)

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

print("CUDA:", torch.cuda.is_available(), torch.cuda.get_device_name(0) if torch.cuda.is_available() else None)
print("bf16 supported:", torch.cuda.is_available() and torch.cuda.is_bf16_supported())Offerto da

CUDA: True NVIDIA A100 80GB PCIe
bf16 supported: TrueOfferto da

2. Preprocessing del dataset LaTeX OCR per il training a bassa quantità di dati

In questo tutorial usiamo il dataset LaTeX OCR disponibile su Hugging Face. Ogni esempio nel dataset è composto da:

• un’immagine contenente un’espressione matematica, e
• il corrispondente testo sorgente LaTeX come output target.

Il dataset fornisce più configurazioni. Per mantenere il training efficiente e allineato al nostro obiettivo di apprendere con pochi dati, controlleremo esplicitamente la dimensione del dataset.

Per prima cosa, carichiamo gli split di training e validation. Per simulare uno scenario di fine-tuning con pochi dati, mescoliamo casualmente il dataset e selezioniamo un piccolo sottoinsieme:

• 1.000 campioni per il training
• 200 campioni per la validation

Questo mantiene il training veloce pur consentendo al modello di generalizzare.

DATASET_NAME = "full"
raw_train = load_dataset("linxy/LaTeX_OCR", name=DATASET_NAME, split="train")
raw_val   = load_dataset("linxy/LaTeX_OCR", name=DATASET_NAME, split="validation")

train_ds = raw_train.shuffle(seed=42).select(range(1000))
val_ds   = raw_val.shuffle(seed=42).select(range(200)) 

print(train_ds, val_ds)
print("Columns:", train_ds.column_names)Offerto da

Output:

Dataset({
    features: ['image', 'text'],
    num_rows: 1000
}) Dataset({
    features: ['image', 'text'],
    num_rows: 200
})
Columns: ['image', 'text']Offerto da

Per comprendere meglio i dati, diamo un’occhiata a un singolo esempio di training. Prima, ispeziona l’immagine.

train_ds[10]["image"]Offerto da

Poi esamina la stringa LaTeX grezza associata a quell’immagine.

train_ds[10]["text"]Offerto da

'G ( \\beta , \\tilde { \\mu } ) = \\left( \\frac { \\pi \\mu \\Gamma ( \\frac { \\lambda } { \\lambda + 1 } ) } { 2 \\Gamma ( \\frac { 1 } { \\lambda + 1 } ) } \\right) ^ { \\frac { 1 } { 2 \\lambda } } g _ { 0 } ( \\beta ) g _ { S } ( \\beta , Z ) ,'Offerto da

Infine, possiamo renderizzare questa espressione LaTeX direttamente nel notebook per verificare che il testo corrisponda al contenuto dell’immagine.

from IPython.display import display, Math, Latex

latex = train_ds[10]["text"]
display(Math(latex))Offerto da

3. Inizializzare l’encoder–decoder multimodale T5Gemma-2

Ora che il dataset è pronto, il passo successivo è caricare il modello T5Gemma-2 e il relativo processor. Useremo la variante 270M–270M, che offre un buon equilibrio tra capacità ed efficienza e si adatta comodamente a una GPU consumer.

MODEL_ID = "google/t5gemma-2-270m-270m"

processor = AutoProcessor.from_pretrained(MODEL_ID, use_fast=True)
model = AutoModelForSeq2SeqLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,   # A100 -> bf16
    device_map="auto",
)Offerto da

Il tokenizer è accessibile tramite il processor. Alcuni modelli encoder–decoder non definiscono un token di padding di default, quindi lo aggiungiamo se necessario e ridimensioniamo di conseguenza gli embedding dei token del modello.

tokenizer = processor.tokenizer
if tokenizer.pad_token is None:
    tokenizer.add_special_tokens({"pad_token": tokenizer.eos_token})
    model.resize_token_embeddings(len(tokenizer))Offerto da

Per l’OCR LaTeX guidiamo il modello con un breve prompt testuale abbinato all’immagine di input. Questo prompt specifica chiaramente il compito e vincola il formato dell’output.

Infine, definiamo le lunghezze massime di sequenza per l’input dell’encoder e per il target del decoder. Questi valori sono sufficienti per la maggior parte delle espressioni matematiche, mantenendo sotto controllo l’uso di memoria.

PROMPT = "<start_of_image> Convert this image to LaTeX. Output only LaTeX."
MAX_INPUT_LEN  = 128
MAX_TARGET_LEN = 256Offerto da

4. Eseguire l’inferenza di base: verifica delle prestazioni zero-shot

Prima di iniziare il fine-tuning, è utile eseguire una inferenza di base usando il modello T5Gemma-2 preaddestrato. Questo ci aiuta a capire come il modello si comporta nel compito di OCR LaTeX senza alcun training specifico e ci fornisce un punto di riferimento per i miglioramenti successivi.

Ecco un breve schema del processo:

1. Scegli un’immagine di esempio dal dataset di training.
2. Prepara gli input del modello combinando l’immagine e il prompt testuale nel formato tensoriale richiesto.
3. Esegui l’inferenza in modalità evaluation con il calcolo dei gradienti disabilitato per risparmiare memoria e accelerare la generazione.
4. Usa il beam search con un numero limitato di beam e leggere penalità di ripetizione per un output più pulito.
5. Decodifica i token generati in testo una volta completato il processo.

image = train_ds[20]["image"]

# prepare inputs
model_inputs = processor(text=PROMPT, images=image, return_tensors="pt")
model_inputs = {k: v.to("cuda") for k, v in model_inputs.items()}

# run inference
model = model.eval()
with torch.inference_mode():
    generation = model.generate(
        **model_inputs,
        max_new_tokens=100,
        do_sample=False,
        num_beams=3,    
        repetition_penalty=1.2,
        no_repeat_ngram_size=4,
        early_stopping=True,
    )

# decode
pred = processor.decode(generation[0], skip_special_tokens=True)

print("\n--- Model output ---")
print(pred)Offerto da

Esempio di output dal modello preaddestrato:

--- Model output ---
Use this command to change this image's format:\usepackage[T1]{fontenc}\usepackage[utf8]{inputenc}\DeclareUnicodeUTF8{1234567890123}\begin{document}<tex>$\begin{equation}\begin{aligned}P=&\frac{1+(-)\frac{1}{2}\left(F_{1}(-)\frac{F_{2}(-Offerto da

Ora vediamo il LaTeX ground truth per la stessa immagine.

train_ds[20]["text"]Offerto da

Ground truth:

'{ \\cal P } = \\frac { 1 + ( - ) ^ { F } { \\cal I } _ { 4 } ( - ) ^ { F _ { L } } } { 2 } .'Offerto da

Come previsto, il modello preaddestrato non produce una trascrizione LaTeX corretta. Genera invece boilerplate LaTeX generico e frammenti non correlati. Questo comportamento è normale, poiché il modello base non è stato addestrato specificamente per l’OCR LaTeX.

5. Creare un data collator personalizzato immagine-testo

In questa sezione costruiamo un data collator immagine–testo personalizzato per batchare correttamente immagini, prompt e target LaTeX per il fine-tuning. Poiché T5Gemma-2 è un modello multimodale encoder–decoder, il collator gioca un ruolo critico nell’assicurare che immagini e testo siano allineati correttamente e passati al modello nel formato atteso.

In particolare, il collator:

• Carica immagini da vari formati possibili e le converte in una rappresentazione RGB coerente
• Garantisce che ogni esempio di training contenga esattamente un’immagine, incapsulata nella struttura attesa dal processor
• Aggancia un prompt di istruzione fisso a ogni immagine per definire chiaramente il compito di OCR
• Tokenizza i target LaTeX separatamente con troncamento e padding controllati
• Maschera i token di padding nelle label in modo che non contribuiscano alla loss di training
• Evita il troncamento sul lato input per prevenire disallineamenti immagine–token

Insieme, questi passaggi garantiscono un training stabile, un calcolo corretto della loss e una convergenza affidabile durante il fine-tuning del modello su dati di OCR LaTeX.

from typing import Any, Dict, List
import torch
from PIL import Image as PILImage

tokenizer = processor.tokenizer
tokenizer.padding_side = "right"

if tokenizer.pad_token is None:
    tokenizer.add_special_tokens({"pad_token": tokenizer.eos_token})

pad_id = tokenizer.pad_token_id

PROMPT = "<start_of_image> Convert this equation image to LaTeX. Output only LaTeX."
MAX_TARGET_LEN = 256

def collate_fn(examples: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
    images, prompts, targets = [], [], []

    for ex in examples:
        im = ex["image"]
        if isinstance(im, PILImage.Image):
            im = im.convert("RGB")
        elif isinstance(im, dict) and "path" in im:
            im = PILImage.open(im["path"]).convert("RGB")
        else:
            raise ValueError(f"Unexpected image type: {type(im)}")

        # IMPORTANT: one image per sample -> nested list
        images.append([im])

        prompts.append(PROMPT)
        targets.append(ex["text"])

    # ✅ NO truncation here (prevents image-token mismatch)
    model_inputs = processor(
        text=prompts,
        images=images,
        padding=True,
        truncation=False,
        return_tensors="pt",
    )

    labels = tokenizer(
        targets,
        padding=True,
        truncation=True,
        max_length=MAX_TARGET_LEN,
        return_tensors="pt",
    )["input_ids"]

    labels[labels == pad_id] = -100
    model_inputs["labels"] = labels
    return model_inputsOfferto da

6. Configurazione del training e impostazioni di ottimizzazione

In questa sezione definiamo la configurazione di training e le impostazioni di ottimizzazione per il fine-tuning di T5Gemma-2 sul compito di OCR LaTeX. Il setup è volutamente semplice e ottimizzato per un training veloce su una singola GPU A100, usando un dataset piccolo e una singola epoca. 

Per ridurre l’overhead e velocizzare, disabilitiamo evaluation e salvataggio dei checkpoint e ci concentriamo su forward e backward pass efficienti.

from transformers import Seq2SeqTrainingArguments

args = Seq2SeqTrainingArguments(
    output_dir="t5gemma2-latex-ocr-1k",

    # --- core training ---
    num_train_epochs=1,
    per_device_train_batch_size=8, 
    gradient_accumulation_steps=1,

    learning_rate=1e-4,
    warmup_steps=15,
    lr_scheduler_type="linear",

    # --- precision / speed ---
    bf16=True,    
    fp16=False,
    tf32=True,     

    # --- memory ---
    gradient_checkpointing=True,

    # --- stop extra work (this is the big speed win) ---
    eval_strategy="no",           
    predict_with_generate=False,          
    save_strategy="no",                 
    report_to="none",

    # --- dataloader ---
    dataloader_num_workers=0,         
    remove_unused_columns=False,

    # --- logging ---
    logging_steps=10,
)Offerto da

Infine, inizializziamo il Seq2SeqTrainer. Passiamo esplicitamente il nostro data collator personalizzato così che il trainer possa costruire correttamente batch multimodali che combinano immagini, prompt di istruzioni e sequenze target LaTeX durante il training.

trainer = Seq2SeqTrainer(
    model=model,
    args=args,
    train_dataset=train_ds,
    eval_dataset=val_ds,
    data_collator=collate_fn,
)Offerto da

7. Fine-tuning di T5Gemma-2 su LaTeX OCR

Con tutto pronto, possiamo avviare il fine-tuning del modello. Il training si lancia con una singola chiamata al trainer.

trainer.train()Offerto da

Durante il training, la loss diminuisce gradualmente, indicando che il modello sta imparando a mappare meglio le immagini di equazioni alle rispettive rappresentazioni LaTeX. 

Anche con un dataset ridotto e una sola epoca, il modello inizia ad adattarsi rapidamente al compito di OCR.

8. Valutazione dopo il fine-tuning

Dopo il fine-tuning, rieseguiamo l’inferenza su un campione di training e su un campione di validation per vedere come sono cambiati gli output del modello. 

Rispetto al baseline, che produceva soprattutto boilerplate LaTeX generico, il modello fine-tunato ora genera LaTeX strutturato che corrisponde da vicino alla forma e ai simboli delle equazioni target.

Partiamo dal test sullo stesso esempio di training di prima. L’immagine viene passata al processor, il modello genera i token e poi li decodifichiamo in LaTeX.

# pick a sample
image = train_ds[10]["image"]

# prepare inputs
model_inputs = processor(text=PROMPT, images=image, return_tensors="pt")
model_inputs = {k: v.to("cuda") for k, v in model_inputs.items()}

# run inference
model = model.eval()
with torch.inference_mode():
    generation = model.generate(
        **model_inputs,
        max_new_tokens=100,
        do_sample=False,
        num_beams=3,    
        repetition_penalty=1.2,
        no_repeat_ngram_size=4,
        early_stopping=True,
    )

# decode
pred = processor.decode(generation[0], skip_special_tokens=True)
print(pred)Offerto da

Come puoi vedere, ora otteniamo un output LaTeX corretto, non testo casuale o non correlato. Nell’esempio di training, la previsione è in gran parte allineata al ground truth, con gli errori rimanenti per lo più in frazioni, indici e qualche token fuori posto.

G ( \beta , \tilde { \mu } ) = \left( \frac { \pi \mu \Gamma } { 2 \Gamma } \lambda _ { + 1 } ^ { \lambda } \right) \right) ^ { \frac { 1 } { 9 \Gamma } g _ { 0 } ( \beta ) g _ { S } ( \theta , Z ) , \right) , \qquad G ( \bar { \Offerto da

Successivamente, testiamo il modello su un campione di validation per verificarne la generalizzazione.

image = val_ds[10]["image"]

# prepare inputs
model_inputs = processor(text=PROMPT, images=image, return_tensors="pt")
model_inputs = {k: v.to("cuda") for k, v in model_inputs.items()}

# run inference
model = model.eval()
with torch.inference_mode():
    generation = model.generate(
        **model_inputs,
        max_new_tokens=100,
        do_sample=False,
        num_beams=3,    
        repetition_penalty=1.2,
        no_repeat_ngram_size=4,
        early_stopping=True,
    )

# decode
pred = processor.decode(generation[0], skip_special_tokens=True)
print(pred)Offerto da

Nell’esempio di validation, il modello segue ancora la struttura e i simboli corretti del LaTeX, anche se occasionalmente sbaglia nel posizionamento delle parentesi, nei termini all’interno delle parentesi e nelle espressioni più lunghe.

f ( p , p ^ { \prime } ) = \ln \left\{ \frac { ( p _ { i } - p _ { j } ) ^ { 2 } } { \left( \frac { \psi ( p _ j ) - g ( p ; p _ { s } ) \psi ( \psi _ { i ] } ) \right\} . . . g ( \psi ; \psi ) - g \left( p ; \psi _Offerto da

Rispetto al ground truth, la struttura complessiva è chiaramente allineata e il modello produce un’approssimazione vicina, piuttosto che un output non correlato.

print(val_ds[10]["text"])Offerto da

f ( p , p ^ { \prime } ) = \ln \left\{ \frac { ( p _ { i } - p _ { j } ) ^ { 2 } } { ( p _ { i } + p _ { j } ) ^ { 2 } } \right\} \left[ \psi ( p _ { j } ) - g ( p _ { i } , p _ { j } ) \psi ( p _ { i } ) \right] .Offerto da

Nel complesso, i risultati post fine-tuning mostrano un chiaro miglioramento. Il modello non indovina più template LaTeX generici, ma produce LaTeX simile a equazioni che somiglia molto ai target del dataset, anche con un set di training ridotto e una breve sessione di fine-tuning.

9. Salvataggio e pubblicazione del modello T5Gemma-2 fine-tunato

Una volta completato il training, il primo passo è salvare localmente il modello fine-tunato in modo da poterlo riutilizzare in seguito per l’inferenza.

trainer.save_model()Offerto da

Poi pubblichiamo il modello su Hugging Face Hub così che altri possano accedervi, riutilizzarlo e svilupparci sopra.

trainer.push_to_hub()Offerto da

Ispezionando i file del repository, potresti notare che la configurazione del processor non è sempre inclusa quando si effettua il push del modello tramite il trainer. 

Poiché il processor è necessario per gestire correttamente sia immagini che testo, lo pubblichiamo esplicitamente a parte per garantire che il modello possa essere caricato e usato senza setup aggiuntivo.

processor.push_to_hub(repo_id="kingabzpro/t5gemma2-latex-ocr-1k")Offerto da

Con questo passaggio, il repository contiene tutto il necessario per caricare modello e processor con una singola chiamata. Ora puoi visitare kingabzpro/t5gemma2-latex-ocr-1k su Hugging Face per accedere al modello fine-tunato e iniziare a usarlo per l’OCR LaTeX o ulteriori esperimenti.

10. Caricare il modello per l’inferenza con Pipelines

Ora che il modello fine-tunato è pubblicato su Hugging Face Hub, possiamo caricarlo direttamente per l’inferenza usando le API di pipeline. Questo è il modo più semplice per testare il modello senza gestire manualmente processor, tokenizer o logica di generazione.

Carichiamo il modello dall’Hub e creiamo una pipeline image-text-to-text:

from transformers import pipeline

generator = pipeline(
    "image-text-to-text",
    model="kingabzpro/t5gemma2-latex-ocr-1k",
)Offerto da

Successivamente eseguiamo l’inferenza su un campione di validation usando lo stesso prompt di istruzioni di prima.

generator(
    val_ds[10]["image"],
    text="<start_of_image> Convert this image to LaTeX. Output only LaTeX.",
    generate_kwargs={"do_sample": False, "max_new_tokens": 100},
)Offerto da

Come puoi vedere, l’output è già molto vicino al LaTeX corretto.

[{'input_text': '<start_of_image> Convert this image to LaTeX. Output only LaTeX.',
  'generated_text': '<start_of_image> Convert this image to LaTeX. Output only LaTeX.f ( p , p ^ { \\prime } ) = \\ln \\left\\{ \\begin{array} { \\begin{array} { \\begin{array} { \\begin{array} { \\begin{array} { \\end{array} \\right\\} \\begin{array} { \\begin{array} { \\begin{array} { \\end{array} \\right\\} \\begin{array} { \\begin{array} { \\begin{array} {'}]Offerto da

Proviamo un altro campione di validation e post-processiamo l’output per mantenere solo la stringa LaTeX.

preds = generator(
    val_ds[30]["image"],
    text="<start_of_image> Convert this image to LaTeX. Output only LaTeX.",
    generate_kwargs={"do_sample": False, "max_new_tokens": 100},
)

prompt = preds[0]["input_text"]
gen = preds[0]["generated_text"]

# remove the prompt if the model echoed it
if gen.startswith(prompt):
    gen = gen[len(prompt):]

# remove any leftover special tokens / separators
gen = gen.replace("<start_of_image>", "").strip()
if gen.startswith("."):
    gen = gen[1:].strip()

print("\nCLEAN PREDICTED LaTeX:\n", gen)Offerto da

Questa volta, il risultato è pulito e direttamente utilizzabile come LaTeX:

CLEAN PREDICTED LaTeX:
 T _ { M N } = \left\{ g N \nu \partial _ { M P _ { } } \cdot \left\{ g ^ { N } \nu \partial _ { M P _ { } } \cdot \left\{ g ^ { N } \nu \partial _ { M P _ { } } \cdot \left\{ g ^ { N } \nu \partial _ { M P _ { } } \cdot \left\{ g ^ {Offerto da

A questo punto, il modello può essere usato come un sistema OCR LaTeX pronto all’uso. Puoi distribuirlo dietro un’API, integrarlo in una pipeline di elaborazione documentale o continuare il fine-tuning con più dati per una precisione ancora migliore.

Se riscontri problemi nell’esecuzione del codice sopra, dai un’occhiata al notebook di supporto. Contiene il codice completo con gli output a ogni passaggio per guidarti nel processo.

Conclusione

Quando ho iniziato ad addestrare il modello, l’ho trattato come un qualsiasi large language model con un image encoder. Dopo vari tentativi falliti, ho capito che questo approccio non funziona per i modelli sequence-to-sequence, encoder–decoder. 

Ho dovuto ripensare l’intero setup, incluso il data collator, il trainer e gli argomenti di training, e perfino come viene eseguita l’inferenza.

In questo tutorial abbiamo visto un workflow end-to-end per il fine-tuning di T5Gemma-2 su un compito di OCR LaTeX, partendo dal setup dell’ambiente e l’ispezione del dataset fino al data collator personalizzato, training efficiente e valutazione post-training. 

Usando un dataset ridotto e una singola GPU A100, abbiamo mostrato che un modello multimodale encoder–decoder può imparare rapidamente a generare LaTeX strutturato e significativo a partire da immagini di equazioni. 

Alla fine, il modello fine-tunato è andato ben oltre l’output boilerplate generico e ha prodotto LaTeX simile a equazioni che si avvicina molto al ground truth, dimostrando quanto possa essere accessibile ed efficace il fine-tuning di moderni modelli open per compiti reali di OCR e comprensione dei documenti.

Se cerchi altri esempi pratici di fine-tuning di LLM, ti consiglio di dare un’occhiata al corso Fine-Tuning with Llama 3.