Upload run_cloud_training.py with huggingface_hub

run_cloud_training.py

@@ -135,7 +135,7 @@ def load_and_prepare_dataset(dataset_name, config):
     """
     Load and prepare the dataset for fine-tuning.
     Sort entries by prompt_number as required.
-    NO TOKENIZATION - DATASET IS ALREADY TOKENIZED
+    Handles both pre-tokenized and string content.
     """
     # Use the default dataset path if no specific path is provided
     if dataset_name == "phi4-cognitive-dataset":
@@ -182,9 +182,35 @@ def load_and_prepare_dataset(dataset_name, config):
         if len(dataset) > 0:
             sample = dataset[0]
             logger.info(f"Sample entry structure: {list(sample.keys())}")
-            if 'conversations' in sample:
+            
+            # Check if dataset is pre-tokenized or contains string content
+            is_pre_tokenized = False
+            
+            if 'input_ids' in sample and isinstance(sample['input_ids'], list) and all(isinstance(x, int) for x in sample['input_ids']):
+                logger.info("Dataset appears to be pre-tokenized with input_ids field")
+                is_pre_tokenized = True
+            elif 'conversations' in sample:
                 logger.info(f"Sample conversations structure: {sample['conversations'][:1]}")
                 
+                # Check if conversations contain pre-tokenized data
+                if isinstance(sample['conversations'], list) and len(sample['conversations']) > 0:
+                    conv = sample['conversations'][0]
+                    if isinstance(conv, dict) and 'input_ids' in conv and isinstance(conv['input_ids'], list):
+                        logger.info("Dataset appears to be pre-tokenized in conversations.input_ids")
+                        is_pre_tokenized = True
+                    elif isinstance(conv, dict) and 'content' in conv:
+                        content = conv['content']
+                        if isinstance(content, list) and all(isinstance(x, int) for x in content):
+                            logger.info("Dataset appears to be pre-tokenized in conversations.content")
+                            is_pre_tokenized = True
+                        else:
+                            logger.info("Dataset appears to contain string content that will need tokenization")
+            
+            if is_pre_tokenized:
+                logger.info("Using pre-tokenized dataset - tokenizer will only be used as fallback")
+            else:
+                logger.info("Dataset contains string content - tokenizer will be used")
+                
         return dataset
     
     except Exception as e:
@@ -208,12 +234,12 @@ def tokenize_string(text, tokenizer):
 # Data collator for pre-tokenized dataset
 class PreTokenizedCollator(DataCollatorMixin):
     """
-    Data collator for pre-tokenized datasets.
-    Expects input_ids and labels already tokenized.
+    Data collator that can handle both pre-tokenized datasets and string content.
+    Will tokenize strings if necessary, but logs warnings.
     """
     def __init__(self, pad_token_id=0, tokenizer=None):
         self.pad_token_id = pad_token_id
-        self.tokenizer = tokenizer  # Keep a reference to the tokenizer for debugging only
+        self.tokenizer = tokenizer  # Keep a reference to the tokenizer for fallback tokenization
         
     def __call__(self, features):
         # Print a sample feature to understand structure
@@ -251,26 +277,33 @@ class PreTokenizedCollator(DataCollatorMixin):
                         feature['input_ids'] = conversations
                     
                     # Case 3: If conversations is a list of dicts with 'content' field
-                    # This should be avoided for pre-tokenized datasets
                     elif isinstance(conversations[0], dict) and 'content' in conversations[0]:
                         content = conversations[0]['content']
                         
                         # If content is already a list of integers, use it directly
                         if isinstance(content, list) and all(isinstance(x, int) for x in content):
                             feature['input_ids'] = content
-                        # AVOID TOKENIZATION: Log warning if content is a string
-                        elif isinstance(content, str):
-                            logger.warning("Found string content in pre-tokenized dataset. This should not happen.")
-                            logger.warning("Skipping this example to avoid tokenization.")
+                        # If content is a string, tokenize it with a warning
+                        elif isinstance(content, str) and self.tokenizer:
+                            logger.warning("Found string content in dataset. Tokenizing as fallback.")
+                            feature['input_ids'] = self.tokenizer.encode(content, add_special_tokens=False)
+                        else:
+                            logger.warning(f"Unexpected content format: {type(content)}")
                             continue
+                    
+                    # Case 4: If conversations is a list of strings
+                    elif all(isinstance(x, str) for x in conversations) and self.tokenizer:
+                        # Join all strings and tokenize
+                        logger.warning("Found string conversations in dataset. Tokenizing as fallback.")
+                        full_text = " ".join(conversations)
+                        feature['input_ids'] = self.tokenizer.encode(full_text, add_special_tokens=False)
             
             # Ensure input_ids is a list of integers
             if 'input_ids' in feature:
-                # AVOID TOKENIZATION: Skip string input_ids
-                if isinstance(feature['input_ids'], str):
-                    logger.warning("Found string input_ids in pre-tokenized dataset. This should not happen.")
-                    logger.warning("Skipping this example to avoid tokenization.")
-                    continue
+                # If input_ids is a string, tokenize it
+                if isinstance(feature['input_ids'], str) and self.tokenizer:
+                    logger.warning("Found string input_ids in dataset. Tokenizing as fallback.")
+                    feature['input_ids'] = self.tokenizer.encode(feature['input_ids'], add_special_tokens=False)
                 # If input_ids is not a list, convert it
                 elif not isinstance(feature['input_ids'], list):
                     try:
@@ -569,7 +602,8 @@ def train(config_path, dataset_name, output_dir):
         logger.info("Successfully applied LoRA with standard PEFT")
             
         # No need to format the dataset - it's already pre-tokenized
-        logger.info("Using pre-tokenized dataset - skipping tokenization step")
+        logger.info("Using dataset with flexible tokenization handling")
+        logger.info("Will use pre-tokenized data if available, or tokenize strings as fallback")
         training_dataset = dataset
         
         # Configure reporting backends with fallbacks