mirror of
https://github.com/hiyouga/LlamaFactory.git
synced 2026-03-23 18:53:23 +08:00
30 lines
983 B
Python
30 lines
983 B
Python
from dataclasses import dataclass
|
|
from transformers import AutoProcessor
|
|
|
|
|
|
@dataclass
|
|
class DataCollatorForVis2Seq:
|
|
processor: AutoProcessor
|
|
|
|
def __call__(self, examples):
|
|
texts = []
|
|
images = []
|
|
for example in examples:
|
|
if len(example["images"]) > 1:
|
|
raise ValueError("This collator only supports one image per example")
|
|
messages = example["messages"]
|
|
text = self.processor.tokenizer.apply_chat_template(
|
|
messages, tokenize=False, add_generation_prompt=False
|
|
)
|
|
texts.append(text)
|
|
images.append(example["images"][0])
|
|
|
|
batch = self.processor(text=texts, images=images, return_tensors="pt", padding=True)
|
|
|
|
labels = batch["input_ids"].clone()
|
|
if self.processor.tokenizer.pad_token_id is not None:
|
|
labels[labels == self.processor.tokenizer.pad_token_id] = -100
|
|
batch["labels"] = labels
|
|
|
|
return batch
|