sail_scripts/train/data_loader.py

import os
import glob

def load_from_folder(path: str) -> str:
    """
    Recursively loads text from all supported files in a directory.
    Also supports loading a single file directly.
    Supported extensions: .txt, .md, .py, .js, .html, .css, .json, .csv, .pdf
    """
    if not os.path.exists(path):
        raise FileNotFoundError(f"The path {path} does not exist.")
    
    path = os.path.expanduser(path)
    text_data = []
    files_found = 0
    
    # Case 1: Single File
    if os.path.isfile(path):
        print(f"Loading single file: {path}")
        try:
            with open(path, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
                if content.strip():
                    text_data.append(content)
                    files_found += 1
        except Exception as e:
            print(f"Error reading file {path}: {e}")
            
        if not text_data:
             print("Warning: File was empty or could not be read.")
        return "\n\n".join(text_data)

    # Case 2: Directory
    print(f"Scanning {path} for data...")
    
    extensions = ['*.txt', '*.md', '*.py', '*.js', '*.html', '*.css', '*.json', '*.csv', '*.pdf', 'data-*'] # Added data-* for the user's specific case if they pass folder
    
    text_columns = ['comment_text', 'text', 'body', 'content', 'comment', 'prompt', 'response', 'instruction', 'input', 'output']

    for ext in extensions:
        # Recursive search for each extension
        search_pattern = os.path.join(path, '**', ext)
        for filepath in glob.glob(search_pattern, recursive=True):
            if not os.path.isfile(filepath): continue
            
            try:
                if filepath.endswith('.csv'):
                    import csv
                    # CSV handling
                    with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
                        # Use Sniffer to detect headers potentially, or just assume DictReader works
                        try:
                            reader = csv.DictReader(f)
                            if reader.fieldnames:
                                # Find valid columns
                                valid_cols = [c for c in reader.fieldnames if c in text_columns]
                                if not valid_cols:
                                    # Fallback: look for any column with "text" in it
                                    valid_cols = [c for c in reader.fieldnames if 'text' in c.lower()]
                                
                                if valid_cols:
                                    print(f"Reading CSV {filepath} columns: {valid_cols}")
                                    csv_text = []
                                    for row in reader:
                                        for col in valid_cols:
                                            if row[col] and row[col].strip():
                                                csv_text.append(row[col])
                                    
                                    if csv_text:
                                        text_data.append("\n".join(csv_text))
                                        files_found += 1
                        except Exception as csv_e:
                             print(f"Error parsing CSV {filepath}: {csv_e}")

                elif filepath.endswith('.pdf'):
                    import pypdf
                    # PDF handling
                    print(f"Reading PDF: {filepath}")
                    try:
                        reader = pypdf.PdfReader(filepath)
                        pdf_text = []
                        for page in reader.pages:
                            text = page.extract_text()
                            if text:
                                pdf_text.append(text)
                        
                        if pdf_text:
                            text_data.append("\n".join(pdf_text))
                            files_found += 1
                    except Exception as pdf_e:
                        print(f"Error reading PDF {filepath}: {pdf_e}")

                elif filepath.endswith('.json'):
                    import json
                    # JSON handling - try to read as structured dataset
                    with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
                        try:
                            data = json.load(f)
                            json_text = []
                            # Case 1: List of dicts (common dataset format)
                            if isinstance(data, list):
                                for item in data:
                                    if isinstance(item, dict):
                                        # Extract value from known text text_columns
                                        for key in item:
                                            if key in text_columns or 'text' in key.lower():
                                                val = item[key]
                                                if isinstance(val, str) and val.strip():
                                                    json_text.append(val)
                            # Case 2: Dict (maybe a config or single record)
                            elif isinstance(data, dict):
                                for key, val in data.items():
                                    if (key in text_columns or 'text' in key.lower()) and isinstance(val, str):
                                        json_text.append(val)
                            
                            if json_text:
                                print(f"Structured JSON loaded from {filepath}")
                                text_data.append("\n".join(json_text))
                                files_found += 1
                            else:
                                # Fallback: Just dump the whole thing as string if no structure found
                                text_data.append(json.dumps(data, indent=2))
                                files_found += 1
                        except Exception as json_e:
                            print(f"Error parsing JSON {filepath}: {json_e}")

                else:
                    # Regular text handling (includes .txt, .md, and extensionless if matched by data-*)
                    with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
                        content = f.read()
                        if content.strip(): # Skip empty files
                            text_data.append(content)
                            files_found += 1
            except Exception as e:
                print(f"Error reading {filepath}: {e}")
                
    print(f"Found {files_found} files. Loading {len(text_data)} chunks.")
    return "\n\n".join(text_data)

if __name__ == "__main__":
    # Test
    # print(load_from_folder(".")) 
    pass