starai-llmdatacleaner 0.1.0

use llm to clean data for start ai

LLMDataCleaner

Use Large Language Models to efficiently clean and standardize your data for machine learning projects.

Installation

You can install the package using your preferred package manager:

Using pip

pip install llmdatacleaner-starai

Using uv

uv add llmdatacleaner-starai

Using poetry

poetry add llmdatacleaner-starai

Configuration

LLMDataCleaner requires an OpenAI API key. You can provide it in one of the following ways:

1. Set as an environment variable:

   export OPENAI_API_KEY="your-api-key"
   

2. Use a .env file in your project directory:

   OPENAI_API_KEY="your-api-key"
   

3. Pass it directly when creating the cleaner:

   cleaner = OpenAIDataCleaner(
       openai_api_key="your-api-key",
       data_cleaning_prompt=prompt,
       output_format=format
   )
   

Usage

Basic Example

from llmdatacleaner import OpenAIDataCleaner

# Define your data cleaning prompt
prompt = """Return the cleaned data into json formats.
1. clean the key 'university' into these categories ['985', '211', 'g5', 'other'] 
2. clean the key 'major' into these categories ['CS', 'EE', 'ME', 'other'] 
3. clean the key 'gpa' from different scale to 4.0 scale

data: {data}
"""

# Define your output format using JSON Schema
format = {
    "format": {
        "type": "json_schema",
        "name": "cleaned_data",
        "schema": {
            "type": "object",
            "properties": {
                "university": {"type": "string", "enum": ["985", "211", "g5", "other"]},
                "major": {"type": "string", "enum": ["CS", "EE", "ME", "other"]},
                "gpa": {"type": "number"},
            },
            "required": ["university", "major", "gpa"],
            "additionalProperties": False,
        },
        "strict": True,
    }
}

# Create the data cleaner
cleaner = OpenAIDataCleaner(
    data_cleaning_prompt=prompt, 
    output_format=format,
    verbose=True  # Optional: enable detailed logging
)

# Clean a single data item
item = {"university": "Stanford University", "major": "Computer Science", "gpa": 4.0}
result = cleaner.invoke(item)

print(result)

Synchronous Processing

For more examples of synchronous data processing, check out the example notebook: 1_test.ipynb

Asynchronous Processing

For processing many items concurrently, use the asynchronous interface:

import asyncio
from llmdatacleaner import OpenAIDataCleaner

# Set up your cleaner as shown above
# ...

async def process_data_async(items):
    tasks = [cleaner.ainvoke(item) for item in items]
    results = await asyncio.gather(*tasks)
    return results

# Process multiple items concurrently
items = [
    {"university": "Tsinghua University", "major": "Computer Science", "gpa": 3.8},
    {"university": "Stanford University", "major": "Mechanical Engineering", "gpa": 4.0},
    # Add more items...
]

results = asyncio.run(process_data_async(items))

For a complete asynchronous processing example, see test.py.

Batch Processing

For extremely large datasets, you can use OpenAI's batch API. An example is provided in 2_batch.ipynb.

API Reference

OpenAIDataCleaner

Parameters:

• data_cleaning_prompt (str): The prompt that instructs the LLM how to clean the data.
• output_format (dict): Format specification for the cleaned output.
• openai_api_key (str, optional): OpenAI API key.
• model_name (str, optional): Default is "gpt-4o-mini".
• system_prompt (str, optional): System instruction for the LLM.
• verbose (bool, optional): Whether to enable detailed logging.

Methods:

• invoke(item): Clean a single data item synchronously.
• ainvoke(item): Clean a single data item asynchronously.

Return value (CleaningResult):

{
    "original": item,          # The original data item
    "cleaned": cleaned_data,   # The cleaned data (parsed JSON or raw text)
    "status": "success",       # "success" or "error"
    "error": None              # Error message if status is "error"
}