datasetplus 0.4.4

An enhanced wrapper for Hugging Face datasets with additional functionality

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.