DatasetPlus

中文版本

An enhanced Hugging Face datasets wrapper designed for large model data processing, providing intelligent caching, data augmentation, and filtering capabilities.

🚀 Core Features

1. 🔄 Full Compatibility - Seamless datasets Replacement

• 100% Compatible: Supports all datasets methods and attributes
• Mutual Conversion, Plug and Play: Simply dsp = DatasetPlus(ds) to convert back: ds = dsp.ds
• Static Method Support: All Dataset static methods and class methods can be called directly

2. 🧠 Intelligent Caching System - Zero Loss for Large Model Calls

• Automatic Function-level Caching: Automatically generates cache based on function content, even if interrupted due to network instability, insufficient quota, etc., previous results won't be lost
• Jupyter Friendly: Even if you forget to assign variables, results can be recovered from cache
• Resume from Breakpoint: Supports continuing processing after interruption, automatically reads cache for processed data

3. 📈 Data Augmentation - One Row Becomes Multiple Rows

• Array Auto-expansion: Automatically expands arrays returned by map function into multiple rows of data
• LLM Result Parsing: Easy data augmentation with MyLLMTool

4. 🔍 Intelligent Filtering - Auto-delete When Returning None

• Conditional Filtering: Automatically deletes rows when map function returns None
• LLM Intelligent Filtering: Use large models for complex conditional filtering

5. 🎨 Generate from Scratch - Direct Data Generation from Large Models

• iter Method: Supports generating datasets from scratch
• Flexible Generation: Can generate data in any format and content
• Batch Generation: Supports large-scale data generation with automatic caching and parallel processing

📦 Installation

Install from PyPI

pip install datasetplus

Install from Source

git clone https://github.com/yourusername/datasetplus.git
cd datasetplus
pip install -e .

Dependencies Installation

# Basic dependencies
pip install datasets pandas numpy

# Excel support
pip install openpyxl

# LLM support
pip install openai

🎯 Quick Start

Basic Usage

from datasetplus import DatasetPlus, MyLLMTool

# Load dataset
dsp = DatasetPlus.load_dataset("data.jsonl")

# Fully compatible with datasets - plug and play
ds = dsp.ds  # Now ds has all datasets functionality + DatasetPlus enhancements

🔄 Feature 1: Full datasets Compatibility

# All datasets methods can be used directly
ds = dsp.ds  # Get native dataset object
dsp_shuffled = dsp.shuffle(seed=42)
dsp_split = dsp.train_test_split(test_size=0.2)
dsp_filtered = dsp.filter(lambda x: len(x['text']) > 100)

# pandas can also seamlessly connect
dsp_df = dsp.to_pandas()
dsp = DatasetPlus.from_pandas(dsp_df)

# Static methods are fully supported
dsp_from_dict = DatasetPlus.from_dict({"text": ["hello", "world"]})
dsp_from_hf = DatasetPlus.from_pretrained("squad")

# Seamless switching with native datasets
from datasets import Dataset
ds = Dataset.from_dict({"text": ["test"]})
dsp = DatasetPlus(ds)  # Directly wrap existing dataset

# Jupyter-friendly display
dsp = DatasetPlus.from_dict({"text": ["a", "b", "c"]})
dsp
----------------
DatasetPlus({
    features: ['text'],
    num_rows: 3
})
~~~~~~~~~~~~~~~~~

## Humanized slicing and display logic 1
dsp[0] # Equivalent to ds.select(range(0)), of course dsp also supports dsp.select(range(0))
----------------
{'text': 'a'}

## Humanized slicing and display logic 2
dsp[1:2] # Equivalent to ds.select(range(1,2))
----------------
DatasetPlus({
    features: ['text'],
    num_rows: 1

## Humanized slicing and display logic 3
dsp[1:]
----------------
DatasetPlus({
    features: ['text'],
    num_rows: 2
})

🧠 Feature 2: Intelligent Caching - Zero Loss for Large Model Calls

# Define processing function containing large model calls
def enhance_with_llm(example):
    # Initialize LLM tool (needs to be instantiated internally for multiprocessing)
    llm = MyLLMTool(
        model_name="gpt-3.5-turbo",
        base_url="https://api.openai.com/v1",
        api_key="your-api-key"
    )
    # Call large model for data enhancement
    prompt = f"Please generate a summary for the following text: {example['text']}"
    summary = llm.getResult(prompt)
    example['summary'] = summary
    return example

# First run - will call large model
dsp_enhanced = dsp.map(enhance_with_llm, num_proc=4, cache=True)

# Forgot to assign in Jupyter? No problem!
# Even if you ran: dsp.map(enhance_with_llm, cache=True)  # Forgot assignment
# You can still recover results with:
dsp_enhanced = dsp.map(enhance_with_llm, cache=True)  # Auto-read from cache, won't call LLM again

# Continuing after interruption is also fine, processed data will be automatically skipped

📈 Feature 3: Data Augmentation - One Row Becomes Multiple Rows

def expand_data_with_llm(example):
    # Use LLM to generate multiple related questions
    prompt = f"Generate 3 related questions based on the following text, return in JSON array format: {example['text']}"
    questions_json = llm.getResult(prompt)
    
    try:
        questions = json.loads(questions_json)
        # Return array, DatasetPlus will automatically expand into multiple rows
        return [{
            'original_text': example['text'],
            'question': q,
            'source': 'llm_generated'
        } for q in questions]
    except:
        return example  # Return original data on parsing failure or delete directly: return None

# Original data: 100 rows
# After processing: might become 300 rows (3 questions generated per row)
dsp_expanded = dsp.map(expand_data_with_llm, cache=True)
print(f"Original data: {len(dsp)} rows")
print(f"After expansion: {len(dsp_expanded)} rows")

🔍 Feature 4: Intelligent Filtering - Auto-delete When Returning None

def filter_with_llm(example):
    # Use LLM for quality assessment
    prompt = f"""Evaluate the quality of the following text, return JSON format: {{"quality": "high/mid/low"}}
    Text: {example['text']}"""
    
    result = llm.getResult(prompt)
    try:
        quality_data = json.loads(result)
        quality = quality_data.get('quality', 'low')
        
        # Only keep high-quality data, others returning None will be auto-deleted
        if quality == 'high':
            example['quality_score'] = quality
            return example
        else:
            return None  # Auto-delete low-quality data
    except:
        return None  # Also delete on parsing failure

# Original data: 1000 rows
# After filtering: might only have 200 rows of high-quality data
dsp_filtered = dsp.map(filter_with_llm, cache=True)
print(f"Before filtering: {len(dsp)} rows")
print(f"After filtering: {len(dsp_filtered)} rows")

🎨 Feature 5: Generate from Scratch - Direct Data Generation from Large Models

# Use iter method to generate data directly from large models
def generate_dialogues(example):
    llm = MyLLMTool(
        model_name="gpt-3.5-turbo",
        base_url="https://api.openai.com/v1",
        api_key="your-api-key"
    )
    
    # Prompt: Generate 10 customer service dialogues
    prompt = """Please generate 10 different customer service dialogues, each containing questions and answers.
    Requirements:
    1. Each dialogue has users asking different specific questions
    2. Customer service provides professional answers
    3. Return JSON array format: [{"user": "User question 1", "assistant": "Service answer 1", "category": "Question category 1"}, ...]
    4. Cover different types of questions (technical support, after-sales service, product consultation, etc.)
    """
    
    try:
        result = llm.getResult(prompt)
        dialogues_data = json.loads(result)
        
        # Return array, DatasetPlus will automatically expand into multiple rows
        return [{
            'batch_id': example['id'],
            'dialogue_id': i,
            'user': dialogue.get('user', ''),
            'assistant': dialogue.get('assistant', ''),
            'category': dialogue.get('category', ''),
            'source': 'generated'
        } for i, dialogue in enumerate(dialogues_data)]
    except Exception as e:
        print(f"Generation failed: {e}")
        return None  # Skip on generation failure

# Generate 10 batches of dialogue data, each batch containing 10 dialogues
dsp_generated = DatasetPlus.iter(
    iterate_num=10,           # Generate 10 batches of data
    fn=generate_dialogues,    # Generation function
    num_proc=2,              # 2 processes in parallel
    cache=False              # iter cache defaults to False
)

print(f"Generated {len(dsp_generated)} dialogue data")  # Should be 100 (10 batches × 10 dialogues)
print(dsp_generated[0])  # View first generated data

📁 Supported Data Formats

• JSON/JSONL: Standard JSON and JSON Lines format
• CSV: Comma-separated values files
• Excel: .xlsx and .xls files
• Hugging Face Datasets: Any dataset from the Hub
• Dataframe, datasets: Support pandas DataFrame and Hugging Face datasets
• Directory Batch Loading: Automatically merge multiple files in directory

🔧 Advanced Features

Intelligent Caching Mechanism

# Cache based on function content hash, ensures recalculation when function changes
def process_v1(example):
    return {"result": example["text"].upper()}  # Version 1

def process_v2(example):
    return {"result": example["text"].lower()}  # Version 2

# Different functions generate different caches, no interference
ds1 = ds.map(process_v1, cache=True)  # Cache A
ds2 = ds.map(process_v2, cache=True)  # Cache B

Batch Processing and Multiprocessing

# Efficient processing of large datasets
dsp = DatasetPlus.load_dataset("large_dataset.jsonl")
dsp_processed = dsp.map(
    enhance_with_llm,
    num_proc=8,           # 8 processes in parallel
    max_inner_num=1000,   # Process 1000 items per batch
    cache=True            # Enable caching
)

Directory Batch Loading

# Automatically load and merge all supported files in directory
dsp = DatasetPlus.load_dataset_plus("./data_folder/")
# Supports mixed formats: data_folder/
#   ├── file1.jsonl
#   ├── file2.csv
#   └── file3.xlsx

Professional Excel Processing

from datasetplus import DatasetPlusExcels

# Professional Excel file processing
excel_dsp = DatasetPlusExcels("spreadsheet.xlsx")

# Support multi-sheet processing
sheet_names = excel_dsp.get_sheet_names()
for sheet in sheet_names:
    sheet_data = excel_dsp.get_sheet_data(sheet)
    dsp_processed = excel_dsp.map(lambda x: {'cleaned': x['column'].strip()})

📚 API Reference

DatasetPlus

Enhanced dataset processing class, fully compatible with Hugging Face datasets.

Core Methods

• map(fn, num_proc=1, max_inner_num=1000, cache=True): Enhanced mapping function
  • fn: Processing function, supports returning arrays (auto-expand) and None (auto-delete)
  • cache: Intelligent caching, automatically generates cache keys based on function content
  • num_proc: Multiprocess parallel processing
  • max_inner_num: Batch processing size

Static Methods

• load_dataset(file_name, output_file): Load single dataset file
• load_dataset_plus(input_path, output_file): Load from file, directory, or Hub
• from_pandas(df): Create from pandas DataFrame
• from_dict(data): Create from dictionary
• from_pretrained(path): Load from Hugging Face Hub
• iter(iterate_num, fn, num_proc=1, max_inner_num=1000, cache=True): Generate from scratch, iteratively generate data
  • iterate_num: Number of data to generate
  • fn: Generation function, receives example with id, returns generated data
  • num_proc: Multiprocess parallel processing
  • cache: Enable caching, avoid duplicate generation

Compatibility

# All datasets methods can be used directly
ds.shuffle()          # Shuffle data
ds.filter()           # Filter data
ds.select()           # Select data
ds.train_test_split() # Split data
ds.save_to_disk()     # Save to disk
# ... and all other datasets methods

MyLLMTool

Large model calling tool, supports OpenAI-compatible APIs.

Initialization

llm = MyLLMTool(
    model_name="gpt-3.5-turbo",      # Model name
    base_url="https://api.openai.com/v1",  # API base URL
    api_key="your-api-key"           # API key
)

Methods

• getResult(query, sys_prompt=None, temperature=0.7, top_p=1, max_tokens=2048, model_name=None): Get LLM response
  • query: User query
  • sys_prompt: System prompt
  • temperature: Temperature parameter
  • max_tokens: Maximum tokens

DatasetPlusExcels

Professional Excel file processing class.

Methods

• __init__(file_path, output_file): Initialize Excel processor
• get_sheet_names(): Get all sheet names
• get_sheet_data(sheet_name): Get data from specified sheet

🎯 Real-world Use Cases

Case 1: Large-scale Data Annotation

# Use LLM for sentiment analysis annotation on large amounts of text
def sentiment_labeling(example):
    prompt = f"Analyze the sentiment of the following text, return positive/negative/neutral: {example['text']}"
    sentiment = llm.getResult(prompt)
    example['sentiment'] = sentiment.strip()
    return example

# Process 100,000 data points, supports resume from breakpoint
dsp_labeled = dsp.map(sentiment_labeling, cache=True, num_proc=4)

Case 2: Data Quality Filtering

# Use LLM to filter high-quality training data
def quality_filter(example):
    prompt = f"Rate text quality (1-5 points): {example['text']}"
    score = llm.getResult(prompt)
    try:
        if int(score) >= 4:
            return example
        else:
            return None  # Auto-delete low-quality data
    except:
        return None

dsp_filtered = dsp.map(quality_filter, cache=True)

Case 3: Data Augmentation

# Generate multiple variants for each data point
def data_augmentation(example):
    prompt = f"Generate 3 synonymous rewrites for the following text: {example['text']}"
    variants = llm.getResult(prompt).split('\n')
    
    # Return array, automatically expand into multiple rows
    return [{
        'text': variant.strip(),
        'label': example['label'],
        'source': 'augmented'
    } for variant in variants if variant.strip()]

dsp_augmented = dsp.map(data_augmentation, cache=True)

Case 4: Generate Training Data from Scratch

# Use LLM to generate training data from scratch
def generate_qa_pairs(example):
    llm = MyLLMTool(
        model_name="gpt-3.5-turbo",
        base_url="https://api.openai.com/v1",
        api_key="your-api-key"
    )
    
    # Prompt for generating Q&A pairs
    prompt = """Generate a Q&A pair about Python programming.
    Requirements:
    1. Question should be specific and practical
    2. Answer should be accurate and detailed
    3. Return JSON format: {"question": "question", "answer": "answer", "difficulty": "easy/medium/hard"}
    """
    
    try:
        result = llm.getResult(prompt)
        qa_data = json.loads(result)
        return {
            'id': example['id'],
            'question': qa_data.get('question', ''),
            'answer': qa_data.get('answer', ''),
            'difficulty': qa_data.get('difficulty', 'medium'),
            'domain': 'python_programming',
            'generated_at': datetime.now().isoformat()
        }
    except Exception as e:
        print(f"Generation failed: {e}")
        return None

# Generate 1000 Python programming Q&A pairs
dsp_qa_dataset = DatasetPlus.iter(
    iterate_num=1000,
    fn=generate_qa_pairs,
    num_proc=4,
    cache=True
)

print(f"Successfully generated {len(dsp_qa_dataset)} Q&A pairs")

💡 Best Practices

1. Caching Strategy

• Always enable caching: cache=True
• Large model call friendly, even if interrupted due to network instability, insufficient quota, etc., previous results won't be lost
• Automatically recalculates when function is modified

2. Performance Optimization

• Set num_proc reasonably (based on maximum concurrency the large model can accept)
• Adjust max_inner_num (maximum memory data storage, writes to disk for persistence every max_inner_num)
• Use batch processing for large datasets

3. Error Handling

def robust_processing(example):
    try:
        # LLM call
        result = llm.getResult(prompt)
        return process_result(result)
    except Exception as e:
        print(f"Processing failed: {e}")
        return None  # Failed data auto-deleted

📋 System Requirements

• Python: >= 3.7
• datasets: >= 2.0.0
• pandas: >= 1.3.0
• numpy: >= 1.21.0
• openpyxl: >= 3.0.0 (Excel support)
• openai: >= 1.0.0 (LLM support)

🤝 Contributing

Pull Requests are welcome! Please ensure:

• Code follows project standards
• Add appropriate tests
• Update relevant documentation

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

📝 Changelog

v0.2.0 (Latest)

• ✨ Added intelligent caching system
• ✨ Support array auto-expansion
• ✨ Support None auto-filtering
• ✨ Full datasets API compatibility
• ✨ Added MyLLMTool large model tool

v0.1.0

• 🎉 Initial release
• 📁 Basic dataset loading functionality
• 📊 Excel file support
• ⚡ Caching and batch processing
• 📂 Directory loading support

🛠️ Auxiliary Tools

DatasetPlus provides two powerful auxiliary tools to enhance your data processing workflow:

MyLLMTool

A comprehensive LLM (Large Language Model) calling tool that supports OpenAI-compatible APIs.

Key Features

• Multi-model Support: Compatible with OpenAI and other OpenAI-compatible APIs
• Flexible Configuration: Customizable parameters for different use cases
• Easy Integration: Seamlessly integrates with DatasetPlus processing workflows

Methods

• __init__(model_name, base_url, api_key): Initialize the LLM tool

  • model_name: The name of the model to use (e.g., "gpt-3.5-turbo")
  • base_url: API base URL (e.g., "https://api.openai.com/v1")
  • api_key: Your API key for authentication

• getResult(query, sys_prompt=None, temperature=0.7, top_p=1, max_tokens=2048, model_name=None): Get LLM response

  • query: User query or prompt
  • sys_prompt: System prompt to set context (optional)
  • temperature: Controls randomness (0.0-2.0)
  • top_p: Controls diversity via nucleus sampling
  • max_tokens: Maximum number of tokens to generate
  • model_name: Override the default model for this request

Usage Example

from datasetplus import MyLLMTool

llm = MyLLMTool(
    model_name="gpt-3.5-turbo",
    base_url="https://api.openai.com/v1",
    api_key="your-api-key"
)

result = llm.getResult(
    query="Summarize this text: ...",
    sys_prompt="You are a helpful assistant.",
    temperature=0.7
)

DataTool

A utility class providing various data processing and parsing functions for common data manipulation tasks.

Key Features

• JSON Parsing: Safe JSON extraction and parsing from text
• File Operations: Read and sample data from files
• Data Validation: Check data structure compliance
• Format Conversion: Convert between different data formats

Methods

• parse_json_safe(text_str): Extract and parse JSON objects/arrays from text

  • text_str: Input string that may contain embedded JSON
  • Returns: List of parsed Python objects (dicts or lists)

• get_prompt(file_path): Read text file and return content as string

  • file_path: Path to the text file
  • Returns: File content as concatenated string

• check(row): Validate data structure for message format

  • row: Data row to validate
  • Returns: Boolean indicating if structure is valid

• check_with_system(row): Check if data has system message format

  • row: Data row to validate
  • Returns: Boolean indicating if has valid system message

• parse_messages(str_row): Parse message format from string

  • str_row: String containing message data
  • Returns: Parsed message object or None

• parse_json(str, json_tag=False): Parse JSON with error handling

  • str: JSON string to parse
  • json_tag: Whether to extract from json code blocks
  • Returns: Parsed object or None on failure

• sample_from_file(file_path, num=-1): Sample lines from text file

  • file_path: Path to the file
  • num: Number of samples (-1 for all)
  • Returns: List of sampled lines

• sample_from(path, num=-1, granularity="auto", exclude=[]): Sample data from files/directories

  • path: File or directory path
  • num: Number of samples (-1 for all)
  • granularity: Sampling granularity ("auto", "file", "line")
  • exclude: Patterns to exclude
  • Returns: List of sampled content

• jsonl2json(source_path, des_path): Convert JSONL to JSON format

  • source_path: Source JSONL file path
  • des_path: Destination JSON file path

Usage Example

from datasetplus import DataTool

# Parse JSON from text
json_data = DataTool.parse_json_safe('Some text {"key": "value"} more text')

# Read prompt from file
prompt = DataTool.get_prompt('prompt.txt')

# Validate data structure
is_valid = DataTool.check(data_row)

# Sample from file
samples = DataTool.sample_from_file('data.txt', num=10)

These auxiliary tools are designed to work seamlessly with DatasetPlus workflows, providing essential utilities for LLM integration and data processing tasks.