Repository: wrmedford/llm
Pull request #1: (feat) fix interleaving strategy and create preprocessing plugin structure
Author: login: wrmedford
Created: 2025-04-20T16:45:24Z | Merged: 2025-04-20T16:49:22Z



Diff:
diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md
index cc17a0f..9e9ee7b 100644
--- a/docs/CONFIGURATION.md
+++ b/docs/CONFIGURATION.md
@@ -104,9 +104,19 @@ dataset_config:
       split: "train"                      # Dataset split to use
       streaming: true                     # Load data progressively (recommended for large datasets)
       weight: 0.7                         # Sampling weight for interleaving
-      # --- Text Extraction (Choose ONE) ---
-      text_field: "text"                  # Field containing the raw text
-      # text_template: "Instruction:\n{instruction}\n\nInput:\n{input}\n\nOutput:\n{output}" # OR: Format string using fields from the dataset
+      # --- Text Extraction (Specify ONE, or leave blank for auto-detection) ---
+      # Option 1: Use a format string with fields from the dataset
+      # text_template: "Instruction:\n{instruction}\n\nInput:\n{input}\n\nOutput:\n{output}"
+      # Option 2: Specify the single field containing the raw text
+      text_field: "text"
+      # Option 3: Specify a preprocessing function to apply to a single field's content
+      # preprocessing_function_name: "format_reasoning" # Name matching a key in PREPROCESSING_FUNCTIONS in datasets.py
+      # text_field: "raw_data"                         # The field to apply the processor to (required if using preprocessing)
+      # --- Priority Order ---
+      # 1. If 'preprocessing_function_name' is set, it's applied to 'text_field', and the result becomes the final 'text'.
+      # 2. Else if 'text_template' is set, it's used to format fields into 'text'.
+      # 3. Else if 'text_field' is set, its content is used as 'text'.
+      # 4. Otherwise, auto-detection is attempted using common field names ('text', 'content', etc.).
 
     - name: "c4"
       path: "allenai/c4"
diff --git a/llm/data/datasets.py b/llm/data/datasets.py
index d563fb1..91fecf7 100644
--- a/llm/data/datasets.py
+++ b/llm/data/datasets.py
@@ -12,13 +12,42 @@
 import os
 import logging
 import os
-from typing import Dict, List
+from typing import Dict, List, Any
 
 import tiktoken # Import tiktoken
 from datasets import IterableDataset, interleave_datasets, load_dataset
 # from transformers import PreTrainedTokenizerBase # No longer needed
 
 
+# --- Preprocessing Functions ---
+
+def format_reasoning_steps(data: Any) -> str:
+    """
+    Example processor: Concatenates 'question' and 'answer' fields
+    if data is a dictionary, otherwise returns data as string.
+    Handles potential nested structures or simple strings.
+    """
+    if isinstance(data, dict):
+        question = data.get("question", "")
+        answer = data.get("answer", "")
+        # Simple concatenation example
+        return f"Question: {question}\nAnswer: {answer}"
+    elif isinstance(data, str):
+        # If the field already contains a string, return it directly
+        return data
+    else:
+        # Handle other unexpected types gracefully
+        return str(data) if data is not None else ""
+
+PREPROCESSING_FUNCTIONS = {
+    "format_reasoning": format_reasoning_steps,
+    # Add more functions here as needed, e.g.:
+    # "extract_code_blocks": extract_code_blocks_function,
+}
+
+# --- End Preprocessing Functions ---
+
+
 def interleave_datasets_with_weights(
     datasets: List[IterableDataset], weights: List[float]
 ) -> IterableDataset:
@@ -56,6 +85,7 @@ def prepare_datasets(config: Dict) -> IterableDataset:
     weights = []
 
     for dataset_info in config["datasets"]:
+        dataset_name = dataset_info.get("name", dataset_info.get("path", "unknown")) # Get a name for logging
         # Load the dataset
         try:
             dataset_path = dataset_info["path"]
@@ -85,81 +115,162 @@ def prepare_datasets(config: Dict) -> IterableDataset:
                     streaming=dataset_info.get("streaming", True),
                 )
 
-            # Map to standardize the text field name
-            text_field = dataset_info.get("text_field", "text")
+            # Determine how to extract text: template, specific field, or auto-detect
+            text_template = dataset_info.get("text_template")
+            text_field = dataset_info.get("text_field")
+            extraction_mode = "auto" # Default
+            final_text_field = None # Field used if mode is 'field' or 'auto'
 
-            # First, check the dataset schema to see what fields are available
+            # Get available fields from a sample
+            available_fields = []
             try:
-                # Get a sample to examine the schema
                 sample = next(iter(dataset.take(1)))
                 available_fields = list(sample.keys())
-                
-                # Log available fields for debugging
-                logger.info(f"Dataset {dataset_info['name']} has fields: {available_fields}")
-                
-                # Verify if the specified text field exists
+                logger.info(f"Dataset {dataset_info['name']} - Available fields: {available_fields}")
+            except StopIteration:
+                logger.warning(f"Dataset {dataset_info['name']} appears to be empty. Skipping.")
+                continue
+            except Exception as e:
+                logger.warning(f"Could not get sample from dataset {dataset_name}: {e}. Proceeding with caution.")
+                # Assume standard fields might exist if we can't get a sample
+
+            # --- Determine Text Extraction/Processing Method ---
+            preprocessing_function_name = dataset_info.get("preprocessing_function_name")
+            text_template = dataset_info.get("text_template")
+            text_field = dataset_info.get("text_field")
+            processed_via_map = False # Flag to track if dataset was processed by a map function
+
+            if preprocessing_function_name:
+                if not text_field:
+                    logger.error(
+                        f"Dataset {dataset_name} specifies 'preprocessing_function_name' "
+                        f"but not 'text_field'. Preprocessing requires a source field. Skipping dataset."
+                    )
+                    continue
                 if text_field not in available_fields:
-                    # Try to find alternative text fields
-                    text_field_candidates = ["text", "content", "document", "passage", "context", "input", "instruction", "query", "question", "code"]
-                    
-                    # Find the first matching field
-                    alternative_field = None
-                    for field in text_field_candidates:
-                        if field in available_fields:
-                            alternative_field = field
-                            break
-                    
-                    # If no standard field found, look for string fields with substantial content
-                    if alternative_field is None:
-                        for field, value in sample.items():
-                            if isinstance(value, str) and len(value) > 100:  # More conservative threshold
-                                alternative_field = field
-                                break
-                    
-                    if alternative_field:
-                        logger.warning(
-                            f"Specified text field '{text_field}' not found in dataset {dataset_info['name']}. "
-                            f"Using '{alternative_field}' instead. Please update your configuration."
-                        )
-                        text_field = alternative_field
-                    else:
+                    logger.error(
+                        f"Specified text_field '{text_field}' for preprocessing not found in dataset {dataset_name}. "
+                        f"Available fields: {available_fields}. Skipping dataset."
+                    )
+                    continue
+                if preprocessing_function_name not in PREPROCESSING_FUNCTIONS:
+                    logger.error(
+                        f"Unknown preprocessing_function_name '{preprocessing_function_name}' for dataset {dataset_name}. "
+                        f"Available functions: {list(PREPROCESSING_FUNCTIONS.keys())}. Skipping dataset."
+                    )
+                    continue
+
+                logger.info(f"Applying preprocessing function '{preprocessing_function_name}' to field '{text_field}' for dataset {dataset_name}")
+                processor = PREPROCESSING_FUNCTIONS[preprocessing_function_name]
+
+                def preprocess_map_fn(example):
+                    try:
+                        source_data = example.get(text_field)
+                        # Pass the entire example to the processor? Or just the field?
+                        # Let's assume processor takes just the field content for now.
+                        if source_data is None:
+                            logger.warning(f"Field '{text_field}' is None in an example from {dataset_name}. Skipping preprocessing for this example.")
+                            return {"text": ""} # Mark for filtering
+                        processed_text = processor(source_data) # Apply the function
+                        # Ensure output is string
+                        return {"text": str(processed_text) if processed_text is not None else ""}
+                    except Exception as e:
+                        logger.error(f"Error applying preprocessing function '{preprocessing_function_name}' to example in {dataset_name}: {e}")
+                        return {"text": ""} # Mark for filtering
+
+                # Apply the preprocessing map function
+                dataset = dataset.map(preprocess_map_fn, batched=False) # Process individually for robustness
+                processed_via_map = True
+                # The 'text' field is now ready, skip further template/field extraction logic
+
+            elif text_template:
+                # Validate template fields exist in the sample (if sample available)
+                if sample: # Only validate if sample is available
+                    try:
+                        # Basic check: does the template string format with the sample?
+                        _ = text_template.format(**sample)
+                        logger.info(f"Using text_template for dataset {dataset_name}")
+                    except KeyError as e:
                         logger.error(
-                            f"Could not find suitable text field in dataset {dataset_info['name']}. "
-                            f"Available fields: {available_fields}. Please specify the correct text_field in your configuration."
+                            f"Field {e} required by text_template not found in dataset {dataset_name}. "
+                            f"Available fields: {available_fields}. Skipping dataset."
                         )
-                        # Skip this dataset instead of using empty strings
                         continue
-            except Exception as e:
-                logger.warning(f"Error examining dataset schema for {dataset_info['name']}: {e}")
-                # Continue with the specified text_field and hope for the best
-
-            def map_fn(example):
-                # Prioritize the specified text_field
-                if text_field in example and isinstance(example[text_field], str):
-                    return {"text": example[text_field]}
-
-                # Fallback logic (simplified) - find the first non-empty string field
-                for k, v in example.items():
-                    if isinstance(v, str) and v.strip():
-                        if k != text_field: # Log only if using a fallback field
-                             logger.warning(
-                                 f"Specified text field '{text_field}' not found or invalid. Using fallback field '{k}' for dataset {dataset_info['name']}."
-                             )
-                        return {"text": v}
+                    except Exception as e:
+                        logger.warning(
+                            f"Error validating text_template for dataset {dataset_name}: {e}. Proceeding with caution."
+                        )
+                else:
+                     logger.info(f"Using text_template for dataset {dataset_name} (sample unavailable for validation).")
+
+                def template_map_fn(example):
+                    try:
+                        text = text_template.format(**example)
+                        return {"text": str(text) if text is not None else ""}
+                    except KeyError as e:
+                        logger.warning(f"Missing field {e} in example from {dataset_name} when applying template. Skipping example.")
+                        return {"text": ""}
+                    except Exception as e:
+                        logger.error(f"Error applying template to example in {dataset_name}: {e}")
+                        return {"text": ""}
+                dataset = dataset.map(template_map_fn, batched=False)
+                processed_via_map = True
+
+            elif text_field:
+                if text_field in available_fields:
+                    logger.info(f"Using specified text_field '{text_field}' for dataset {dataset_name}")
+                    # Check if the field is already 'text', if so, no mapping needed unless type conversion is required
+                    if text_field == "text" and sample and isinstance(sample.get("text"), str):
+                         logger.info(f"Field '{text_field}' is already the target 'text' field and is a string. No mapping applied.")
+                    else:
+                        # Map to rename/convert field to 'text'
+                        def field_map_fn(example):
+                            text = example.get(text_field)
+                            return {"text": str(text) if text is not None else ""}
+                        dataset = dataset.map(field_map_fn, batched=False)
+                        processed_via_map = True
                 else:
-                    # If no suitable field found, log error and return empty text
-                    # This will create essentially empty examples that should be filtered out
                     logger.error(
-                        f"No suitable text content found in example from dataset {dataset_info['name']}. "
-                        f"Fields: {list(example.keys())}"
+                        f"Specified text_field '{text_field}' not found in dataset {dataset_name}. "
+                        f"Available fields: {available_fields}. Skipping dataset."
                     )
-                    # If no suitable string field found, return empty string for filtering
-                    return {"text": ""}
-
-            # Map to standardize text field
-            dataset = dataset.map(map_fn)
+                    continue
+            else:
+                # Auto-detect: Try common field names
+                logger.info(f"No preprocessing, template, or text_field specified for {dataset_name}. Attempting auto-detection.")
+                detected_field = None
+                common_fields = ["text", "content", "document", "passage", "code", "instruction", "input", "output", "prompt", "completion"]
+                if sample: # Only auto-detect if sample is available
+                    for field in common_fields:
+                        if field in available_fields:
+                            # Check if the field contains string data in the sample
+                            if isinstance(sample.get(field), str):
+                                detected_field = field
+                                logger.warning(
+                                    f"Auto-detected text field '{detected_field}' for dataset {dataset_name}. "
+                                    "Consider specifying 'text_field' or 'text_template' in config for clarity."
+                                )
+                                break
+                            else:
+                                logger.warning(f"Field '{field}' found but is not a string in sample. Skipping this candidate.")
+
+                if detected_field:
+                    # Map if detected field is not 'text'
+                    if detected_field != "text":
+                        def auto_detect_map_fn(example):
+                            text = example.get(detected_field)
+                            return {"text": str(text) if text is not None else ""}
+                        dataset = dataset.map(auto_detect_map_fn, batched=False)
+                        processed_via_map = True
+                else:
+                    logger.error(
+                        f"Could not auto-detect a suitable text field for dataset {dataset_name}. "
+                        f"Available fields: {available_fields}. Please specify 'text_field' or 'text_template'. Skipping dataset."
+                    )
+                    continue
 
-            # Filter out examples with empty or whitespace-only text
+            # --- Common Post-processing ---
+            # Filter out examples with empty or whitespace-only text from the 'text' field
             def filter_empty(example):
                 # Check if 'text' exists and is a string before stripping
                 text_content = example.get("text")
@@ -167,24 +278,20 @@ def filter_empty(example):
 
             filtered_dataset = dataset.filter(filter_empty)
 
-            # Log dataset readiness (logging counts might be difficult/slow with streaming)
-            try:
-                # This might not work for all streaming datasets
-                logger.info(f"Dataset {dataset_info['name']} ready for training")
-            except Exception:
-                pass  # Skip count logging for streaming datasets
-            
+            # Log dataset readiness
+            logger.info(f"Dataset {dataset_name} processed and ready for interleaving.")
+
             datasets.append(filtered_dataset)
             weights.append(dataset_info.get("weight", 1.0))
 
             logger.info(
-                f"Added dataset {dataset_info['name']} with weight "
+                f"Added dataset {dataset_name} with weight "
                 f"{dataset_info.get('weight', 1.0)}"
             )
 
         except Exception as e:
-            logger.error(f"Error loading dataset {dataset_info['name']}: {e}")
-            logger.error(f"Skipping dataset {dataset_info['name']}")
+            logger.exception(f"Error processing dataset {dataset_name}: {e}") # Use logger.exception to include traceback
+            logger.error(f"Skipping dataset {dataset_name}")
             continue
 
     if not datasets:


Reviews:
user: login: gemini-code-assist[bot]
body: Hello @wrmedford, I'm Gemini Code Assist[^1]! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

## Summary of Changes
This pull request focuses on improving the data loading and preprocessing pipeline. It introduces a more flexible way to specify how text is extracted from datasets, including the option to use preprocessing functions. The changes also aim to improve the robustness of the data loading process by handling potential errors and providing more informative logging.

### Highlights
* **Interleaving Strategy Fix**: The pull request addresses and fixes issues related to the interleaving strategy of datasets, ensuring a more balanced and reliable data loading process.
* **Preprocessing Plugin Structure**: A new structure for preprocessing plugins is introduced, allowing users to define custom functions to transform the text data before it's used for training. This includes adding a dictionary of available preprocessing functions.
* **Flexible Text Extraction**: The configuration now supports specifying a preprocessing function, a text template, or a specific field for extracting text from datasets, with a clear priority order. Auto-detection is used as a fallback.
* **Improved Error Handling and Logging**: The pull request includes enhanced error handling and logging to provide more informative messages when datasets are loaded and processed, making it easier to debug potential issues.

### Changelog
<details>
<summary>Click here to see the changelog</summary>

* **docs/CONFIGURATION.md**
  * Updated the documentation to reflect the new text extraction options, including preprocessing functions, text templates, and specific fields.
  * Added a priority order to explain how the text extraction method is determined.
* **llm/data/datasets.py**
  * Introduced a dictionary `PREPROCESSING_FUNCTIONS` to store available preprocessing functions.
  * Added the `format_reasoning_steps` function as an example preprocessing function.
  * Modified the `prepare_datasets` function to support preprocessing functions, text templates, and specific fields for text extraction.
  * Implemented error handling and logging to improve the robustness of the data loading process.
  * Added auto-detection of text fields as a fallback option.
  * Implemented filtering of empty or whitespace-only text examples.
  * Improved logging messages to provide more information about the data loading and processing steps.

</details>


<details>

<summary><b>Using Gemini Code Assist</b></summary>
<br>

The full guide for Gemini Code Assist can be found on our [documentation page](https://developers.google.com/gemini-code-assist/docs/review-github-code), here are some quick tips.

<b>Invoking Gemini</b>

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either `/gemini <command>` or `@gemini-code-assist <command>`. Below is a summary of the supported commands.

Feature | Command | Description
--- | --- | ---
Code Review | `/gemini review` | Performs a code review for the current pull request in its current state.
Pull Request Summary | `/gemini summary` | Provides a summary of the current pull request in its current state.
Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments.
Help | `/gemini help` | Displays a list of available commands.

<b>Customization</b>

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a `.gemini/` folder in the base of the repository. Detailed instructions can be found [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github).

<b>Limitations & Feedback</b>

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with :thumbsup: and :thumbsdown: on @gemini-code-assist comments to provide feedback.

</details>

<b>You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the [Gemini Code Assist IDE Extension](https://cloud.google.com/products/gemini/code-assist).</b>


[^1]: Review the [Privacy Notices](https://policies.google.com/privacy), [Generative AI Prohibited Use Policy](https://policies.google.com/terms/generative-ai/use-policy), [Terms of Service](https://policies.google.com/terms), and learn how to configure Gemini Code Assist in GitHub [here](https://developers.google.com/gemini-code-assist/docs/review-github-code). Gemini can make mistakes, so double check it and [use code with caution](https://support.google.com/legal/answer/13505487).

---

> A dataset's deep,
> Text tangled, secrets sleep.
> Processors awake.


state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-20T16:45:32Z
user: login: gemini-code-assist[bot]
body: ## Code Review
This pull request introduces a preprocessing plugin structure and fixes the interleaving strategy for datasets. The changes are well-structured and address a clear need for more flexible data handling. The addition of preprocessing functions and the refined logic for text extraction are valuable improvements.

### Summary of Findings
* **Error Handling in Preprocessing**: The error handling in the `preprocess_map_fn` function could be improved by logging the specific exception that occurs during preprocessing. This would provide more detailed information for debugging.
* **Clarity in Text Extraction Logic**: The logic for determining the text extraction method is complex. Consider adding more comments or restructuring the code to improve readability and maintainability.
* **Logging Consistency**: Ensure consistent use of logging levels (info, warning, error) throughout the code to provide informative feedback during dataset processing.

### Merge Readiness
The pull request is almost ready for merging. Addressing the error handling in preprocessing and improving the clarity of the text extraction logic would significantly enhance the quality of the code. I am unable to directly approve this pull request, and recommend that others review and approve this code before merging. At a minimum, the high severity issues should be addressed before merging.
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-20T16:45:40Z
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This pull request aims to fix issues with the dataset interleaving strategy and introduce a flexible preprocessing plugin structure for text extraction.  
- Introduces a new preprocessing function (format_reasoning_steps) and a PREPROCESSING_FUNCTIONS registry.  
- Refactors dataset handling to support multiple text extraction modes including field-based, template-based, preprocessing function, and auto-detection.  
- Updates configuration documentation to clearly explain the available text extraction options.

### Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

| File                   | Description                                                            |
| ---------------------- | ---------------------------------------------------------------------- |
| llm/data/datasets.py   | Adds preprocessing functions and refactors text extraction logic.    |
| docs/CONFIGURATION.md  | Updates text extraction configuration documentation for clarity.     |


<details>
<summary>Comments suppressed due to low confidence (1)</summary>

**llm/data/datasets.py:129**
* Consider using the already assigned variable 'dataset_name' instead of 'dataset_info["name"]' to maintain consistency in logging messages.
```
logger.info(f"Dataset {dataset_info['name']} has fields: {available_fields}")
```
</details>


state: COMMENTED
author_association: NONE
submitted_at: 2025-04-20T16:46:05Z

Review comments:
diff_hunk: @@ -85,106 +115,183 @@ def prepare_datasets(config: Dict) -> IterableDataset:
                     streaming=dataset_info.get("streaming", True),
                 )
 
-            # Map to standardize the text field name
-            text_field = dataset_info.get("text_field", "text")
+            # Determine how to extract text: template, specific field, or auto-detect
+            text_template = dataset_info.get("text_template")
+            text_field = dataset_info.get("text_field")
+            extraction_mode = "auto" # Default
+            final_text_field = None # Field used if mode is 'field' or 'auto'
 
-            # First, check the dataset schema to see what fields are available
+            # Get available fields from a sample
+            available_fields = []
             try:
-                # Get a sample to examine the schema
                 sample = next(iter(dataset.take(1)))
                 available_fields = list(sample.keys())
-                
-                # Log available fields for debugging
-                logger.info(f"Dataset {dataset_info['name']} has fields: {available_fields}")
-                
-                # Verify if the specified text field exists
+                logger.info(f"Dataset {dataset_info['name']} - Available fields: {available_fields}")
+            except StopIteration:
+                logger.warning(f"Dataset {dataset_info['name']} appears to be empty. Skipping.")
+                continue
+            except Exception as e:
+                logger.warning(f"Could not get sample from dataset {dataset_name}: {e}. Proceeding with caution.")
+                # Assume standard fields might exist if we can't get a sample
+
+            # --- Determine Text Extraction/Processing Method ---
+            preprocessing_function_name = dataset_info.get("preprocessing_function_name")
+            text_template = dataset_info.get("text_template")
+            text_field = dataset_info.get("text_field")
+            processed_via_map = False # Flag to track if dataset was processed by a map function
+
+            if preprocessing_function_name:
+                if not text_field:
+                    logger.error(
+                        f"Dataset {dataset_name} specifies 'preprocessing_function_name' "
+                        f"but not 'text_field'. Preprocessing requires a source field. Skipping dataset."
+                    )
+                    continue
                 if text_field not in available_fields:
-                    # Try to find alternative text fields
-                    text_field_candidates = ["text", "content", "document", "passage", "context", "input", "instruction", "query", "question", "code"]
-                    
-                    # Find the first matching field
-                    alternative_field = None
-                    for field in text_field_candidates:
-                        if field in available_fields:
-                            alternative_field = field
-                            break
-                    
-                    # If no standard field found, look for string fields with substantial content
-                    if alternative_field is None:
-                        for field, value in sample.items():
-                            if isinstance(value, str) and len(value) > 100:  # More conservative threshold
-                                alternative_field = field
-                                break
-                    
-                    if alternative_field:
-                        logger.warning(
-                            f"Specified text field '{text_field}' not found in dataset {dataset_info['name']}. "
-                            f"Using '{alternative_field}' instead. Please update your configuration."
-                        )
-                        text_field = alternative_field
-                    else:
+                    logger.error(
+                        f"Specified text_field '{text_field}' for preprocessing not found in dataset {dataset_name}. "
+                        f"Available fields: {available_fields}. Skipping dataset."
+                    )
+                    continue
+                if preprocessing_function_name not in PREPROCESSING_FUNCTIONS:
+                    logger.error(
+                        f"Unknown preprocessing_function_name '{preprocessing_function_name}' for dataset {dataset_name}. "
+                        f"Available functions: {list(PREPROCESSING_FUNCTIONS.keys())}. Skipping dataset."
+                    )
+                    continue
+
+                logger.info(f"Applying preprocessing function '{preprocessing_function_name}' to field '{text_field}' for dataset {dataset_name}")
+                processor = PREPROCESSING_FUNCTIONS[preprocessing_function_name]
+
+                def preprocess_map_fn(example):
+                    try:
+                        source_data = example.get(text_field)
+                        # Pass the entire example to the processor? Or just the field?
+                        # Let's assume processor takes just the field content for now.
+                        if source_data is None:
+                            logger.warning(f"Field '{text_field}' is None in an example from {dataset_name}. Skipping preprocessing for this example.")
+                            return {"text": ""} # Mark for filtering
+                        processed_text = processor(source_data) # Apply the function
+                        # Ensure output is string
+                        return {"text": str(processed_text) if processed_text is not None else ""}
+                    except Exception as e:
+                        logger.error(f"Error applying preprocessing function '{preprocessing_function_name}' to example in {dataset_name}: {e}")
+                        return {"text": ""} # Mark for filtering
+
+                # Apply the preprocessing map function
+                dataset = dataset.map(preprocess_map_fn, batched=False) # Process individually for robustness
+                processed_via_map = True
+                # The 'text' field is now ready, skip further template/field extraction logic
+
+            elif text_template:
+                # Validate template fields exist in the sample (if sample available)
+                if sample: # Only validate if sample is available
+                    try:
+                        # Basic check: does the template string format with the sample?
+                        _ = text_template.format(**sample)
+                        logger.info(f"Using text_template for dataset {dataset_name}")
+                    except KeyError as e:
                         logger.error(
-                            f"Could not find suitable text field in dataset {dataset_info['name']}. "
-                            f"Available fields: {available_fields}. Please specify the correct text_field in your configuration."
+                            f"Field {e} required by text_template not found in dataset {dataset_name}. "
+                            f"Available fields: {available_fields}. Skipping dataset."
                         )
-                        # Skip this dataset instead of using empty strings
                         continue
-            except Exception as e:
-                logger.warning(f"Error examining dataset schema for {dataset_info['name']}: {e}")
-                # Continue with the specified text_field and hope for the best
-
-            def map_fn(example):
-                # Prioritize the specified text_field
-                if text_field in example and isinstance(example[text_field], str):
-                    return {"text": example[text_field]}
-
-                # Fallback logic (simplified) - find the first non-empty string field
-                for k, v in example.items():
-                    if isinstance(v, str) and v.strip():
-                        if k != text_field: # Log only if using a fallback field
-                             logger.warning(
-                                 f"Specified text field '{text_field}' not found or invalid. Using fallback field '{k}' for dataset {dataset_info['name']}."
-                             )
-                        return {"text": v}
+                    except Exception as e:
+                        logger.warning(
+                            f"Error validating text_template for dataset {dataset_name}: {e}. Proceeding with caution."
+                        )
+                else:
+                     logger.info(f"Using text_template for dataset {dataset_name} (sample unavailable for validation).")
+
+                def template_map_fn(example):
+                    try:
+                        text = text_template.format(**example)
+                        return {"text": str(text) if text is not None else ""}
+                    except KeyError as e:
+                        logger.warning(f"Missing field {e} in example from {dataset_name} when applying template. Skipping example.")
+                        return {"text": ""}
+                    except Exception as e:
+                        logger.error(f"Error applying template to example in {dataset_name}: {e}")
+                        return {"text": ""}
+                dataset = dataset.map(template_map_fn, batched=False)
+                processed_via_map = True
+
+            elif text_field:
+                if text_field in available_fields:
+                    logger.info(f"Using specified text_field '{text_field}' for dataset {dataset_name}")
+                    # Check if the field is already 'text', if so, no mapping needed unless type conversion is required
+                    if text_field == "text" and sample and isinstance(sample.get("text"), str):
+                         logger.info(f"Field '{text_field}' is already the target 'text' field and is a string. No mapping applied.")
+                    else:
+                        # Map to rename/convert field to 'text'
+                        def field_map_fn(example):
+                            text = example.get(text_field)
+                            return {"text": str(text) if text is not None else ""}
+                        dataset = dataset.map(field_map_fn, batched=False)
+                        processed_via_map = True
                 else:
-                    # If no suitable field found, log error and return empty text
-                    # This will create essentially empty examples that should be filtered out
                     logger.error(
-                        f"No suitable text content found in example from dataset {dataset_info['name']}. "
-                        f"Fields: {list(example.keys())}"
+                        f"Specified text_field '{text_field}' not found in dataset {dataset_name}. "
+                        f"Available fields: {available_fields}. Skipping dataset."
                     )
-                    # If no suitable string field found, return empty string for filtering
-                    return {"text": ""}
-
-            # Map to standardize text field
-            dataset = dataset.map(map_fn)
+                    continue
+            else:
+                # Auto-detect: Try common field names
+                logger.info(f"No preprocessing, template, or text_field specified for {dataset_name}. Attempting auto-detection.")
+                detected_field = None
+                common_fields = ["text", "content", "document", "passage", "code", "instruction", "input", "output", "prompt", "completion"]
+                if sample: # Only auto-detect if sample is available
+                    for field in common_fields:
+                        if field in available_fields:
+                            # Check if the field contains string data in the sample
+                            if isinstance(sample.get(field), str):
+                                detected_field = field
+                                logger.warning(
+                                    f"Auto-detected text field '{detected_field}' for dataset {dataset_name}. "
+                                    "Consider specifying 'text_field' or 'text_template' in config for clarity."
+                                )
+                                break
+                            else:
+                                logger.warning(f"Field '{field}' found but is not a string in sample. Skipping this candidate.")
+
+                if detected_field:
+                    # Map if detected field is not 'text'
+                    if detected_field != "text":
+                        def auto_detect_map_fn(example):
+                            text = example.get(detected_field)
+                            return {"text": str(text) if text is not None else ""}
+                        dataset = dataset.map(auto_detect_map_fn, batched=False)
+                        processed_via_map = True
+                else:
+                    logger.error(
+                        f"Could not auto-detect a suitable text field for dataset {dataset_name}. "
+                        f"Available fields: {available_fields}. Please specify 'text_field' or 'text_template'. Skipping dataset."
+                    )
+                    continue
 
-            # Filter out examples with empty or whitespace-only text
+            # --- Common Post-processing ---
+            # Filter out examples with empty or whitespace-only text from the 'text' field
             def filter_empty(example):
                 # Check if 'text' exists and is a string before stripping
                 text_content = example.get("text")
                 return isinstance(text_content, str) and len(text_content.strip()) > 0
 
             filtered_dataset = dataset.filter(filter_empty)
 
-            # Log dataset readiness (logging counts might be difficult/slow with streaming)
-            try:
-                # This might not work for all streaming datasets
-                logger.info(f"Dataset {dataset_info['name']} ready for training")
-            except Exception:
-                pass  # Skip count logging for streaming datasets
-            
+            # Log dataset readiness
+            logger.info(f"Dataset {dataset_name} processed and ready for interleaving.")
+
             datasets.append(filtered_dataset)
             weights.append(dataset_info.get("weight", 1.0))
 
             logger.info(
-                f"Added dataset {dataset_info['name']} with weight "
+                f"Added dataset {dataset_name} with weight "
                 f"{dataset_info.get('weight', 1.0)}"
             )
 
         except Exception as e:
-            logger.error(f"Error loading dataset {dataset_info['name']}: {e}")
-            logger.error(f"Skipping dataset {dataset_info['name']}")
+            logger.exception(f"Error processing dataset {dataset_name}: {e}") # Use logger.exception to include traceback
+            logger.error(f"Skipping dataset {dataset_name}")
path: llm/data/datasets.py
user: login: gemini-code-assist[bot]
body: ![high](https://img.shields.io/badge/priority-high-orange)

Using `logger.exception` is great for including the traceback, but consider also including the specific exception message in the main error message for easier readability in logs.

```suggestion
            logger.exception(f"Error processing dataset {dataset_name}: {type(e).__name__} - {e}") # Use logger.exception to include traceback
            logger.error(f"Skipping dataset {dataset_name}")
```
created_at: 2025-04-20T16:45:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -85,106 +115,183 @@ def prepare_datasets(config: Dict) -> IterableDataset:
                     streaming=dataset_info.get("streaming", True),
                 )
 
-            # Map to standardize the text field name
-            text_field = dataset_info.get("text_field", "text")
+            # Determine how to extract text: template, specific field, or auto-detect
+            text_template = dataset_info.get("text_template")
+            text_field = dataset_info.get("text_field")
+            extraction_mode = "auto" # Default
+            final_text_field = None # Field used if mode is 'field' or 'auto'
 
-            # First, check the dataset schema to see what fields are available
+            # Get available fields from a sample
+            available_fields = []
             try:
-                # Get a sample to examine the schema
                 sample = next(iter(dataset.take(1)))
                 available_fields = list(sample.keys())
-                
-                # Log available fields for debugging
-                logger.info(f"Dataset {dataset_info['name']} has fields: {available_fields}")
-                
-                # Verify if the specified text field exists
+                logger.info(f"Dataset {dataset_info['name']} - Available fields: {available_fields}")
+            except StopIteration:
+                logger.warning(f"Dataset {dataset_info['name']} appears to be empty. Skipping.")
+                continue
+            except Exception as e:
+                logger.warning(f"Could not get sample from dataset {dataset_name}: {e}. Proceeding with caution.")
+                # Assume standard fields might exist if we can't get a sample
+
+            # --- Determine Text Extraction/Processing Method ---
+            preprocessing_function_name = dataset_info.get("preprocessing_function_name")
+            text_template = dataset_info.get("text_template")
+            text_field = dataset_info.get("text_field")
+            processed_via_map = False # Flag to track if dataset was processed by a map function
+
+            if preprocessing_function_name:
+                if not text_field:
+                    logger.error(
+                        f"Dataset {dataset_name} specifies 'preprocessing_function_name' "
+                        f"but not 'text_field'. Preprocessing requires a source field. Skipping dataset."
+                    )
+                    continue
                 if text_field not in available_fields:
-                    # Try to find alternative text fields
-                    text_field_candidates = ["text", "content", "document", "passage", "context", "input", "instruction", "query", "question", "code"]
-                    
-                    # Find the first matching field
-                    alternative_field = None
-                    for field in text_field_candidates:
-                        if field in available_fields:
-                            alternative_field = field
-                            break
-                    
-                    # If no standard field found, look for string fields with substantial content
-                    if alternative_field is None:
-                        for field, value in sample.items():
-                            if isinstance(value, str) and len(value) > 100:  # More conservative threshold
-                                alternative_field = field
-                                break
-                    
-                    if alternative_field:
-                        logger.warning(
-                            f"Specified text field '{text_field}' not found in dataset {dataset_info['name']}. "
-                            f"Using '{alternative_field}' instead. Please update your configuration."
-                        )
-                        text_field = alternative_field
-                    else:
+                    logger.error(
+                        f"Specified text_field '{text_field}' for preprocessing not found in dataset {dataset_name}. "
+                        f"Available fields: {available_fields}. Skipping dataset."
+                    )
+                    continue
+                if preprocessing_function_name not in PREPROCESSING_FUNCTIONS:
+                    logger.error(
+                        f"Unknown preprocessing_function_name '{preprocessing_function_name}' for dataset {dataset_name}. "
+                        f"Available functions: {list(PREPROCESSING_FUNCTIONS.keys())}. Skipping dataset."
+                    )
+                    continue
+
+                logger.info(f"Applying preprocessing function '{preprocessing_function_name}' to field '{text_field}' for dataset {dataset_name}")
+                processor = PREPROCESSING_FUNCTIONS[preprocessing_function_name]
+
+                def preprocess_map_fn(example):
+                    try:
+                        source_data = example.get(text_field)
+                        # Pass the entire example to the processor? Or just the field?
+                        # Let's assume processor takes just the field content for now.
+                        if source_data is None:
+                            logger.warning(f"Field '{text_field}' is None in an example from {dataset_name}. Skipping preprocessing for this example.")
+                            return {"text": ""} # Mark for filtering
+                        processed_text = processor(source_data) # Apply the function
+                        # Ensure output is string
+                        return {"text": str(processed_text) if processed_text is not None else ""}
+                    except Exception as e:
+                        logger.error(f"Error applying preprocessing function '{preprocessing_function_name}' to example in {dataset_name}: {e}")
+                        return {"text": ""} # Mark for filtering
path: llm/data/datasets.py
user: login: gemini-code-assist[bot]
body: ![medium](https://img.shields.io/badge/priority-medium-yellow)

Consider logging the specific exception `e` in the error message to provide more detailed information for debugging. This can help identify the root cause of preprocessing failures.

```suggestion
                        logger.error(f"Error applying preprocessing function '{preprocessing_function_name}' to example in {dataset_name}: {e}", exc_info=True)
                        return {"text": ""} # Mark for filtering
```
created_at: 2025-04-20T16:45:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -85,106 +115,183 @@ def prepare_datasets(config: Dict) -> IterableDataset:
                     streaming=dataset_info.get("streaming", True),
                 )
 
-            # Map to standardize the text field name
-            text_field = dataset_info.get("text_field", "text")
+            # Determine how to extract text: template, specific field, or auto-detect
+            text_template = dataset_info.get("text_template")
+            text_field = dataset_info.get("text_field")
+            extraction_mode = "auto" # Default
+            final_text_field = None # Field used if mode is 'field' or 'auto'
 
-            # First, check the dataset schema to see what fields are available
+            # Get available fields from a sample
+            available_fields = []
             try:
-                # Get a sample to examine the schema
                 sample = next(iter(dataset.take(1)))
                 available_fields = list(sample.keys())
-                
-                # Log available fields for debugging
-                logger.info(f"Dataset {dataset_info['name']} has fields: {available_fields}")
-                
-                # Verify if the specified text field exists
+                logger.info(f"Dataset {dataset_info['name']} - Available fields: {available_fields}")
+            except StopIteration:
+                logger.warning(f"Dataset {dataset_info['name']} appears to be empty. Skipping.")
+                continue
+            except Exception as e:
+                logger.warning(f"Could not get sample from dataset {dataset_name}: {e}. Proceeding with caution.")
+                # Assume standard fields might exist if we can't get a sample
+
+            # --- Determine Text Extraction/Processing Method ---
+            preprocessing_function_name = dataset_info.get("preprocessing_function_name")
+            text_template = dataset_info.get("text_template")
+            text_field = dataset_info.get("text_field")
path: llm/data/datasets.py
user: login: Copilot
body: [nitpick] The variables 'text_template' and 'text_field' are being retrieved again later in the preprocessing branch; consider consolidating these assignments to avoid redundancy and potential confusion.
```suggestion

```
created_at: 2025-04-20T16:46:04Z
author_association: NONE

PR comments:
user: login: wrmedford
created_at: 2025-04-20T16:45:36Z
body: /gemini
author_association: OWNER
user: login: gemini-code-assist[bot]
created_at: 2025-04-20T16:45:45Z
body: It looks like you ran the `/gemini` command. To trigger a code review, please use `/gemini review`. For a list of available commands, use `/gemini help`.
author_association: CONTRIBUTOR