Upload scripts/_sft_classes.py with huggingface_hub

scripts/_sft_classes.py

@@ -0,0 +1,97 @@
+
+import json
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+
+
+class PhysicsCoTDataset(Dataset):
+    """Dataset for Qwen2.5-VL SFT with physics CoT."""
+
+    def __init__(self, data_path, processor, max_length=4096):
+        self.processor = processor
+        self.max_length = max_length
+        with open(data_path, 'r', encoding='utf-8') as f:
+            self.records = [json.loads(line) for line in f]
+        print(f"Loaded {len(self.records)} records from {data_path}")
+
+    def __len__(self):
+        return len(self.records)
+
+    def __getitem__(self, idx):
+        record = self.records[idx]
+        messages = record['messages']
+        user_msg = messages[0]
+        image_path = None
+        text_content = ""
+        for content in user_msg['content']:
+            if content['type'] == 'image':
+                image_path = content['image'].replace('file://', '')
+            elif content['type'] == 'text':
+                text_content = content['text']
+        assistant_msg = messages[1]
+        assistant_text = assistant_msg['content'][0]['text']
+        image = Image.open(image_path).convert('RGB')
+        MIN_DIM = 56
+        w, h = image.size
+        if w < MIN_DIM or h < MIN_DIM:
+            scale = max(MIN_DIM / w, MIN_DIM / h)
+            new_w = int(w * scale)
+            new_h = int(h * scale)
+            image = image.resize((new_w, new_h), Image.LANCZOS)
+            if new_w < MIN_DIM or new_h < MIN_DIM:
+                padded = Image.new('RGB', (max(new_w, MIN_DIM), max(new_h, MIN_DIM)), (255, 255, 255))
+                padded.paste(image, (0, 0))
+                image = padded
+        conversation = [
+            {"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": text_content}]},
+            {"role": "assistant", "content": [{"type": "text", "text": assistant_text}]},
+        ]
+        text = self.processor.apply_chat_template(conversation, tokenize=False, add_generation_prompt=False)
+        inputs = self.processor(text=[text], images=[image], padding=False, truncation=True, max_length=self.max_length, return_tensors="pt")
+        input_ids = inputs['input_ids'].squeeze(0)
+        attention_mask = inputs['attention_mask'].squeeze(0)
+        labels = input_ids.clone()
+        assistant_token_str = "<|im_start|>assistant\n"
+        assistant_token_ids = self.processor.tokenizer.encode(assistant_token_str, add_special_tokens=False)
+        input_ids_list = input_ids.tolist()
+        assistant_start = -1
+        for i in range(len(input_ids_list) - len(assistant_token_ids) + 1):
+            if input_ids_list[i:i + len(assistant_token_ids)] == assistant_token_ids:
+                assistant_start = i + len(assistant_token_ids)
+                break
+        if assistant_start > 0:
+            labels[:assistant_start] = -100
+        else:
+            raise ValueError(f"FATAL: assistant start token not found in sample {idx}.")
+        labels[attention_mask == 0] = -100
+        return {
+            'input_ids': input_ids,
+            'attention_mask': attention_mask,
+            'labels': labels,
+            'pixel_values': inputs.get('pixel_values', torch.tensor([])).squeeze(0) if 'pixel_values' in inputs else None,
+            'image_grid_thw': inputs.get('image_grid_thw', torch.tensor([])).squeeze(0) if 'image_grid_thw' in inputs else None,
+        }
+
+
+class VLMDataCollator:
+    """Custom data collator for variable-length VLM inputs."""
+    def __init__(self, processor):
+        self.processor = processor
+        self.pad_token_id = processor.tokenizer.pad_token_id or processor.tokenizer.eos_token_id
+
+    def __call__(self, features):
+        max_len = max(f['input_ids'].size(0) for f in features)
+        input_ids, attention_mask, labels, pixel_values, image_grid_thw = [], [], [], [], []
+        for f in features:
+            seq_len = f['input_ids'].size(0)
+            pad_len = max_len - seq_len
+            input_ids.append(torch.cat([f['input_ids'], torch.full((pad_len,), self.pad_token_id, dtype=f['input_ids'].dtype)]))
+            attention_mask.append(torch.cat([f['attention_mask'], torch.zeros(pad_len, dtype=f['attention_mask'].dtype)]))
+            labels.append(torch.cat([f['labels'], torch.full((pad_len,), -100, dtype=f['labels'].dtype)]))
+            if f.get('pixel_values') is not None: pixel_values.append(f['pixel_values'])
+            if f.get('image_grid_thw') is not None: image_grid_thw.append(f['image_grid_thw'])
+        batch = {'input_ids': torch.stack(input_ids), 'attention_mask': torch.stack(attention_mask), 'labels': torch.stack(labels)}
+        if pixel_values: batch['pixel_values'] = torch.cat(pixel_values, dim=0)
+        if image_grid_thw: batch['image_grid_thw'] = torch.stack(image_grid_thw)
+        return batch