161 lines
4.6 KiB
Python
161 lines
4.6 KiB
Python
"""
|
|
Data collators for dataset processing.
|
|
|
|
This module contains custom data collators for training,
|
|
particularly for VLM/OCR processing.
|
|
"""
|
|
|
|
import torch
|
|
from dataclasses import dataclass
|
|
from typing import Any, List, Optional, Union
|
|
|
|
|
|
@dataclass
|
|
class DeepSeekOCRDataCollator:
|
|
"""
|
|
Data collator for DeepSeek OCR VLM training.
|
|
|
|
Handles:
|
|
- Image processing via processor
|
|
- Text tokenization
|
|
- Proper label masking for instruction fine-tuning
|
|
"""
|
|
processor: Any # Qwen2VLProcessor or similar
|
|
max_length: int = 2048
|
|
ignore_index: int = -100
|
|
|
|
def __call__(self, batch: List[dict]) -> dict:
|
|
"""
|
|
Collate a batch of samples.
|
|
|
|
Args:
|
|
batch: List of dicts, each with 'messages' containing
|
|
[{'role': 'user', 'content': [...]}, {'role': 'assistant', 'content': [...]}]
|
|
|
|
Returns:
|
|
dict with input_ids, attention_mask, labels, pixel_values, etc.
|
|
"""
|
|
from PIL import Image
|
|
|
|
# Extract messages and images
|
|
all_messages = []
|
|
all_images = []
|
|
|
|
for sample in batch:
|
|
messages = sample["messages"]
|
|
all_messages.append(messages)
|
|
|
|
# Extract PIL images from content
|
|
for msg in messages:
|
|
content = msg.get("content", [])
|
|
if isinstance(content, list):
|
|
for item in content:
|
|
if isinstance(item, dict) and item.get("type") == "image":
|
|
img = item.get("image")
|
|
if img is not None and hasattr(img, 'size'): # PIL Image
|
|
all_images.append(img)
|
|
|
|
# Process with the VL processor
|
|
try:
|
|
# Qwen2VL style processing
|
|
texts = [
|
|
self.processor.apply_chat_template(
|
|
msgs, tokenize=False, add_generation_prompt=False
|
|
)
|
|
for msgs in all_messages
|
|
]
|
|
|
|
# Process with images
|
|
inputs = self.processor(
|
|
text=texts,
|
|
images=all_images if all_images else None,
|
|
return_tensors="pt",
|
|
padding=True,
|
|
truncation=True,
|
|
max_length=self.max_length,
|
|
)
|
|
|
|
# Create labels (mask input, keep output)
|
|
labels = inputs["input_ids"].clone()
|
|
|
|
# Simple masking: mask padding tokens
|
|
labels[labels == self.processor.tokenizer.pad_token_id] = self.ignore_index
|
|
|
|
inputs["labels"] = labels
|
|
|
|
return inputs
|
|
|
|
except Exception as e:
|
|
print(f"⚠️ DeepSeekOCRDataCollator error: {e}")
|
|
raise
|
|
|
|
|
|
@dataclass
|
|
class VLMDataCollator:
|
|
"""
|
|
Generic VLM data collator that works with various processors.
|
|
|
|
Supports:
|
|
- Qwen2VL
|
|
- LLaVA
|
|
- Other VL models with compatible processors
|
|
"""
|
|
processor: Any
|
|
max_length: int = 2048
|
|
ignore_index: int = -100
|
|
mask_input_tokens: bool = True # Whether to mask user tokens in labels
|
|
|
|
def __call__(self, batch: List[dict]) -> dict:
|
|
"""
|
|
Collate a batch of VLM samples.
|
|
"""
|
|
all_messages = []
|
|
all_images = []
|
|
|
|
for sample in batch:
|
|
messages = sample.get("messages", [])
|
|
all_messages.append(messages)
|
|
|
|
# Extract images
|
|
for msg in messages:
|
|
content = msg.get("content", [])
|
|
if isinstance(content, list):
|
|
for item in content:
|
|
if isinstance(item, dict):
|
|
img = item.get("image")
|
|
if img is not None:
|
|
all_images.append(img)
|
|
|
|
# Apply chat template
|
|
texts = [
|
|
self.processor.apply_chat_template(
|
|
msgs, tokenize=False, add_generation_prompt=False
|
|
)
|
|
for msgs in all_messages
|
|
]
|
|
|
|
# Process inputs
|
|
inputs = self.processor(
|
|
text=texts,
|
|
images=all_images if all_images else None,
|
|
return_tensors="pt",
|
|
padding=True,
|
|
truncation=True,
|
|
max_length=self.max_length,
|
|
)
|
|
|
|
# Create labels
|
|
labels = inputs["input_ids"].clone()
|
|
|
|
# Mask padding
|
|
if hasattr(self.processor, 'tokenizer'):
|
|
pad_token_id = self.processor.tokenizer.pad_token_id
|
|
else:
|
|
pad_token_id = self.processor.pad_token_id
|
|
|
|
if pad_token_id is not None:
|
|
labels[labels == pad_token_id] = self.ignore_index
|
|
|
|
inputs["labels"] = labels
|
|
|
|
return inputs
|