PaliGemma2 LaTeX OCR 微调实战：公式图片识别与文本差异对比-平芜编程栈

PaliGemma2 LaTeX OCR 微调实战：公式图片识别与文本差异对比

这篇教程是我根据 PaliGemma2 在 LaTeX OCR 任务上的微调复现过程整理出来的。重点演示如何下载公式图片 JSONL 数据集，微调 PaliGemma2-10B，并用可视化 diff、BLEU 和 TER 评估生成公式文本。

LaTeX OCR 的输入是公式图片，输出是 LaTeX 字符串。模型不仅要识别符号，还要生成正确的结构和命令，因此文本差异可视化和序列级指标很重要。

本文会重点跑通以下流程：

下载 LaTeX OCR JSONL 数据集
展示公式图片和目标 LaTeX 文本
使用 QLoRA 微调 PaliGemma2-10B
可视化标注文本和生成文本差异
使用 BLEU 和 TER 评估 OCR 输出

如果你正在系统学习多模态微调、目标检测、OCR 或图像分割，建议收藏本文；配套 notebook、示例图片和运行环境说明后续会继续整理。如果环境配置卡住，可以在评论区说明具体报错。

📚 文章目录

PaliGemma2 LaTeX OCR 微调实战：公式图片识别与文本差异对比
- ⚙️ 环境准备
- 📦 下载 LaTeX OCR 数据集
- 🧾 加载并预览 OCR 数据
- 🧠 加载 PaliGemma2-10B 模型
- 🏋️ 微调 LaTeX OCR 模型
- 🔍 推理并对比生成文本
- 📊 评估 OCR 文本质量
- 📌 小结
- 📚 同系列教程汇总

⚙️ 环境准备

检查 GPU，安装PEFT、bitsandbytes，并安装支持该 notebook 的 transformers 分支。

!nvidia-smi

📦 下载 LaTeX OCR 数据集

从数据集后台获取 unsloth-latex-ocr JSONL 数据集，并查看训练、验证、测试集规模。

!pip install-q supervision peft bitsandbytes

!pip uninstall-y transformers

!pip install-q git+https://github.com/probicheaux/transformers.git@main

fromtypesimportSimpleNamespace# 从数据集后台下载并解压数据集后，修改 DATASET_DIR 指向数据集目录。DATASET_DIR="/content/dataset"# 修改为数据集后台导出的数据集目录dataset=SimpleNamespace(location=DATASET_DIR)

!head-n5{dataset.location}/train/annotations.jsonl

!wc-l<{dataset.location}/train/annotations.jsonl !wc-l<{dataset.location}/valid/annotations.jsonl !wc-l<{dataset.location}/test/annotations.jsonl

# !head -n 10000 {dataset.location}/train/annotations.jsonl > {dataset.location}/train/annotations.sample.jsonl# !head -n 1000 {dataset.location}/valid/annotations.jsonl > {dataset.location}/valid/annotations.sample.jsonl# !head -n 1000 {dataset.location}/test/annotations.jsonl > {dataset.location}/test/annotations.sample.jsonl

🧾 加载并预览 OCR 数据

定义 JSONL 数据集类，并用 HTML 表格展示公式图片和对应 prefix/suffix。

importosimportjsonimportrandomfromPILimportImagefromtorch.utils.dataimportDatasetclassJSONLDataset(Dataset):def__init__(self,jsonl_file_path:str,image_directory_path:str):self.jsonl_file_path=jsonl_file_path self.image_directory_path=image_directory_path self.entries=self._load_entries()def_load_entries(self):entries=[]withopen(self.jsonl_file_path,'r')asfile:forlineinfile:data=json.loads(line)entries.append(data)returnentriesdef__len__(self):returnlen(self.entries)def__getitem__(self,idx:int):ifidx<0oridx>=len(self.entries):raiseIndexError("Index out of range")entry=self.entries[idx]image_path=os.path.join(self.image_directory_path,entry['image'])image=Image.open(image_path)returnimage,entry

train_dataset=JSONLDataset(jsonl_file_path=f"{dataset.location}/train/annotations.jsonl",image_directory_path=f"{dataset.location}/train",)valid_dataset=JSONLDataset(jsonl_file_path=f"{dataset.location}/valid/annotations.jsonl",image_directory_path=f"{dataset.location}/valid",)test_dataset=JSONLDataset(jsonl_file_path=f"{dataset.location}/test/annotations.jsonl",image_directory_path=f"{dataset.location}/test",)

fromIPython.core.displayimportdisplay,HTMLfromPILimportImageimportioimportbase64defpil_image_to_base64(img):"""Convert a PIL image to a base64 string."""buffered=io.BytesIO()img.save(buffered,format="JPEG")img_str=base64.b64encode(buffered.getvalue()).decode("utf-8")returnf"data:image/jpeg;base64,{img_str}"defdisplay_images_and_text(dataset,num_entries=10):""" Display images and their corresponding text side by side in an HTML table. :param dataset: PyTorch dataset to extract images and texts from. :param num_entries: Number of entries to display. """images=[]texts=[]foriinrange(min(num_entries,len(dataset))):img,data=dataset[i]images.append(pil_image_to_base64(img))text=f"Prefix:{data['prefix']}<br>Suffix:{data['suffix']}"texts.append(text)rows=[]forimg,textinzip(images,texts):row_html=f""" <tr> <td><img src="{img}" alt="Image" style="max-width:300px; max-height:300px; object-fit:cover;"></td> <td>{text}</td> </tr> """rows.append(row_html)html_content=f""" <html> <head> <style> body {{ font-family: Arial, sans-serif; margin: 0; padding: 0; }} table {{ width: 100%; border-collapse: collapse; margin: 20px 0; }} th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }} img {{ display: block; margin: auto; }} </style> </head> <body> <table> <tr> <th>Image</th> <th>Text</th> </tr>{''.join(rows)}</table> </body> </html> """display(HTML(html_content))display_images_and_text(train_dataset,num_entries=10)

🧠 加载 PaliGemma2-10B 模型

加载 PaliGemma2-10B processor，并通过 LoRA/QLoRA 控制训练显存。

importtorchfromtransformersimportPaliGemmaProcessor,PaliGemmaForConditionalGeneration MODEL_ID="google/paligemma2-10b-pt-224"DEVICE=torch.device("cuda"iftorch.cuda.is_available()else"cpu")

fromhuggingface_hubimportnotebook_login notebook_login()

processor=PaliGemmaProcessor.from_pretrained(MODEL_ID)

USE_LORA=FalseUSE_QLORA=TrueFREEZE_VISION=False

frompeftimportget_peft_model,LoraConfigfromtransformersimportBitsAndBytesConfigifUSE_LORAorUSE_QLORA:lora_config=LoraConfig(r=8,target_modules=["q_proj","o_proj","k_proj","v_proj","gate_proj","up_proj","down_proj"],task_type="CAUSAL_LM",)ifUSE_QLORA:bnb_config=BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_quant_type="nf4",bnb_4bit_compute_type=torch.bfloat16)model=PaliGemmaForConditionalGeneration.from_pretrained(MODEL_ID,device_map="auto",quantization_config=bnb_configifUSE_QLORAelseNone,torch_dtype=torch.bfloat16)model=get_peft_model(model,lora_config)model=model.to(DEVICE)model.print_trainable_parameters()else:model=PaliGemmaForConditionalGeneration.from_pretrained(MODEL_ID,device_map="auto").to(DEVICE)model=model.to(DEVICE)ifFREEZE_VISION:forparaminmodel.vision_tower.parameters():param.requires_grad=Falseforparaminmodel.multi_modal_projector.parameters():param.requires_grad=FalseTORCH_DTYPE=model.dtype

🏋️ 微调 LaTeX OCR 模型

构造 OCR 任务 collate 函数，使用 Transformers Trainer 训练模型。

fromtransformersimportTrainer,TrainingArgumentsdefcollate_fn(batch):images,labels=zip(*batch)paths=[label["image"]forlabelinlabels]prefixes=["<image>"+label["prefix"]forlabelinlabels]suffixes=[label["suffix"]forlabelinlabels]inputs=processor(text=prefixes,images=images,return_tensors="pt",suffix=suffixes,padding="longest").to(TORCH_DTYPE).to(DEVICE)returninputs args=TrainingArguments(num_train_epochs=3,remove_unused_columns=False,per_device_train_batch_size=3,gradient_accumulation_steps=12,warmup_steps=2,learning_rate=2e-5,weight_decay=1e-6,adam_beta2=0.999,logging_steps=100,optim="paged_adamw_8bit"ifUSE_QLORAelse"adamw_hf",save_strategy="steps",save_steps=1000,save_total_limit=1,output_dir="paligemma2_latex_ocr_v5",bf16=True,report_to=["tensorboard"],dataloader_pin_memory=False)trainer=Trainer(model=model,train_dataset=train_dataset,eval_dataset=valid_dataset,data_collator=collate_fn,args=args)

trainer.train()

🔍 推理并对比生成文本

对测试集样本生成 LaTeX，并用左右对比 diff 标出差异。

# @title Function to render text diffsfromdifflibimportSequenceMatcherfromIPython.core.displayimportdisplay,HTMLdefside_by_side_diff_divs(text1,text2):lines1=text1.splitlines()lines2=text2.splitlines()original_output=[]modified_output=[]forline1,line2inzip(lines1,lines2):words1=line1.split()words2=line2.split()matcher=SequenceMatcher(None,words1,words2)original_line=[]modified_line=[]fortag,i1,i2,j1,j2inmatcher.get_opcodes():iftag=='replace':original_line.append(f"<span class='diff-remove'>{' '.join(words1[i1:i2])}</span>")modified_line.append(f"<span class='diff-add'>{' '.join(words2[j1:j2])}</span>")eliftag=='delete':original_line.append(f"<span class='diff-remove'>{' '.join(words1[i1:i2])}</span>")eliftag=='insert':modified_line.append(f"<span class='diff-add'>{' '.join(words2[j1:j2])}</span>")eliftag=='equal':original_line.append(' '.join(words1[i1:i2]))modified_line.append(' '.join(words2[j1:j2]))original_output.append(' '.join(original_line)+"<br>")modified_output.append(' '.join(modified_line)+"<br>")original_html="<br>"+''.join(original_output)+"<br>"modified_html="<br>"+''.join(modified_output)+"<br>"html=f""" <html> <head> <style> body {{ font-family: Arial, sans-serif; margin: 0; padding: 0; }} .container {{ display: flex; align-items: flex-start; }} .column {{ flex: 1; padding: 10px; white-space: pre-wrap; text-align: left; }} .diff-remove {{ background-color: #d9534f; /* Dark red */ color: white; text-decoration: line-through; border-radius: 4px; padding: 2px 4px; }} .diff-add {{ background-color: #5cb85c; /* Dark green */ color: white; border-radius: 4px; padding: 2px 4px; }} </style> </head> <body> <div class="container"> <div class="column" style="border-right: 1px solid #ccc;">{original_html}</div> <div class="column">{modified_html}</div> </div> </body> </html> """returnhtml

# @title Suffix vs. generated textforiinrange(10):image,label=test_dataset[i]prefix="<image>"+label["prefix"]suffix=label["suffix"]inputs=processor(text=prefix,images=image,return_tensors="pt").to(TORCH_DTYPE).to(DEVICE)prefix_length=inputs["input_ids"].shape[-1]withtorch.inference_mode():generation=model.generate(**inputs,max_new_tokens=256,do_sample=False,num_beams=3)generation=generation[0][prefix_length:]generated_text=processor.decode(generation,skip_special_tokens=True)html_diff=side_by_side_diff_divs(suffix,generated_text)display(image)display(HTML(html_diff))

📊 评估 OCR 文本质量

收集测试集预测结果，使用 BLEU 和 TER 评估生成文本质量。

importnumpyasnp targets=[]predictions=[]withtorch.inference_mode():foriinrange(10):image,label=test_dataset[i]prefix="<image>"+label["prefix"]suffix=label["suffix"]inputs=processor(text=prefix,images=image,return_tensors="pt").to(TORCH_DTYPE).to(DEVICE)prefix_length=inputs["input_ids"].shape[-1]generation=model.generate(**inputs,max_new_tokens=256,do_sample=False,num_beams=3)generation=generation[0][prefix_length:]generated_text=processor.decode(generation,skip_special_tokens=True)targets.append(suffix)predictions.append(generated_text)

!pip install-q evaluate

# @title Calculate BLEUfromevaluateimportload bleu=load("bleu")results=bleu.compute(predictions=predictions,references=targets)print(results)

!pip install-q sacrebleu

# @title Calculate TERfromevaluateimportload ter=load("ter")results=ter.compute(predictions=predictions,references=targets,case_sensitive=True)print(results)