Final cleanup

studio/backend/utils/datasets/data_collators.py

@@ -12,11 +12,10 @@ import torch
 from dataclasses import dataclass
 from typing import Any, List, Optional, Union
 from loggers import get_logger
+
 logger = get_logger(__name__)
 
 
-
-
 @dataclass
 class DataCollatorSpeechSeq2SeqWithPadding:
     """
@@ -26,16 +25,23 @@ class DataCollatorSpeechSeq2SeqWithPadding:
     masks padding in labels with -100, and strips leading BOS token.
     Mirrors the collator from the Whisper.ipynb notebook.
     """
+
     processor: Any
 
     def __call__(self, features: List[dict]) -> dict:
-        input_features = [{"input_features": feature["input_features"]} for feature in features]
-        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
+        input_features = [
+            {"input_features": feature["input_features"]} for feature in features
+        ]
+        batch = self.processor.feature_extractor.pad(
+            input_features, return_tensors = "pt"
+        )
 
         label_features = [{"input_ids": feature["labels"]} for feature in features]
-        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")
+        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors = "pt")
 
-        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)
+        labels = labels_batch["input_ids"].masked_fill(
+            labels_batch.attention_mask.ne(1), -100
+        )
 
         if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
             labels = labels[:, 1:]
@@ -54,6 +60,7 @@ class DeepSeekOCRDataCollator:
     - Text tokenization
     - Proper label masking for instruction fine-tuning
     """
+
     processor: Any  # Qwen2VLProcessor or similar
     max_length: int = 2048
     ignore_index: int = -100
@@ -86,7 +93,7 @@ class DeepSeekOCRDataCollator:
                     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
+                            if img is not None and hasattr(img, "size"):  # PIL Image
                                 all_images.append(img)
 
         # Process with the VL processor
@@ -94,19 +101,19 @@ class DeepSeekOCRDataCollator:
             # Qwen2VL style processing
             texts = [
                 self.processor.apply_chat_template(
-                    msgs, tokenize=False, add_generation_prompt=False
+                    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,
+                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)
@@ -134,6 +141,7 @@ class VLMDataCollator:
     - LLaVA
     - Other VL models with compatible processors
     """
+
     processor: Any
     max_length: int = 2048
     ignore_index: int = -100
@@ -163,26 +171,26 @@ class VLMDataCollator:
         # Apply chat template
         texts = [
             self.processor.apply_chat_template(
-                msgs, tokenize=False, add_generation_prompt=False
+                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,
+            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'):
+        if hasattr(self.processor, "tokenizer"):
             pad_token_id = self.processor.tokenizer.pad_token_id
         else:
             pad_token_id = self.processor.pad_token_id