chatan 0.1.0__tar.gz

chatan-0.1.0/.gitignore

@@ -0,0 +1,10 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv

chatan-0.1.0/.python-version

@@ -0,0 +1 @@
+3.11

chatan-0.1.0/PKG-INFO

@@ -0,0 +1,83 @@
+Metadata-Version: 2.4
+Name: chatan
+Version: 0.1.0
+Summary: Create synthetic datasets with LLM generators and samplers
+Project-URL: Documentation, https://github.com/cdreetz/chatan#readme
+Project-URL: Issues, https://github.com/cdreetz/chatan/issues
+Project-URL: Source, https://github.com/cdreetz/chatan
+Author-email: Christian Reetz <cdreetz@gmail.com>
+License-Expression: MIT
+Keywords: dataset generation,llm,machine learning,synthetic data
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.8
+Requires-Dist: anthropic>=0.7.0
+Requires-Dist: datasets>=2.0.0
+Requires-Dist: numpy>=1.20.0
+Requires-Dist: openai>=1.0.0
+Requires-Dist: pandas>=1.3.0
+Requires-Dist: pydantic>=2.0.0
+Description-Content-Type: text/markdown
+
+## Examples
+
+Prompt a dataset
+
+```
+import chatan
+
+gen = chatan.generator.client("YOUR_OPENAI_API_KEY")
+ds = chatan.dataset("create a QA dataset for finetuning an LLM on pharmacology")
+```
+
+Creating datasets with different data mixes
+
+```
+import uuid
+from chatan import dataset, generator, mix
+
+gen = generator.client("YOUR_OPENAI_API_KEY")
+#generator.client("anthropic", "YOUR_ANTHROPIC_API_KEY")
+
+mix = {
+    "implementation": "Can you implement a matmul kernel in Triton",
+    "conversion": "Convert this pytorch model to Triton",
+    "explanation": "What memory access optimizations are being used here?"
+}
+
+ds = dataset({
+    "id": uuid,
+    "task": sample.choice(mix),
+    "prompt": gen("write a prompt for {task}"),
+    "response": gen("write a response to {prompt}"),
+)}
+```
+
+Augment datasets
+
+```
+import uuid
+from chatan import dataset, generator
+from dataset import load_dataset
+
+gen = generator.client("YOUR_OPENAI_API_KEY")
+hf_dataset = load_dataset("GPU_MODE/KernelBook")
+
+ds = dataset({
+    "id": sample.from_dataset(hf_data, "id", default=sample.uuid()),
+    "prompt": sample.from_dataset(hf_data, "prompt", aug=gen("provide a variation of this prompt")),
+    "response": gen("write a response to {prompt}")
+
+})
+
+```

chatan-0.1.0/README.md

@@ -0,0 +1,52 @@
+## Examples
+
+Prompt a dataset
+
+```
+import chatan
+
+gen = chatan.generator.client("YOUR_OPENAI_API_KEY")
+ds = chatan.dataset("create a QA dataset for finetuning an LLM on pharmacology")
+```
+
+Creating datasets with different data mixes
+
+```
+import uuid
+from chatan import dataset, generator, mix
+
+gen = generator.client("YOUR_OPENAI_API_KEY")
+#generator.client("anthropic", "YOUR_ANTHROPIC_API_KEY")
+
+mix = {
+    "implementation": "Can you implement a matmul kernel in Triton",
+    "conversion": "Convert this pytorch model to Triton",
+    "explanation": "What memory access optimizations are being used here?"
+}
+
+ds = dataset({
+    "id": uuid,
+    "task": sample.choice(mix),
+    "prompt": gen("write a prompt for {task}"),
+    "response": gen("write a response to {prompt}"),
+)}
+```
+
+Augment datasets
+
+```
+import uuid
+from chatan import dataset, generator
+from dataset import load_dataset
+
+gen = generator.client("YOUR_OPENAI_API_KEY")
+hf_dataset = load_dataset("GPU_MODE/KernelBook")
+
+ds = dataset({
+    "id": sample.from_dataset(hf_data, "id", default=sample.uuid()),
+    "prompt": sample.from_dataset(hf_data, "prompt", aug=gen("provide a variation of this prompt")),
+    "response": gen("write a response to {prompt}")
+
+})
+
+```

chatan-0.1.0/main.py

@@ -0,0 +1,21 @@
+import uuid
+from minos import dataset, generator, mix
+
+gen = generator.client("YOUR_OPENAI_API_KEY")
+#generator.client("anthropic", "YOUR_ANTHROPIC_API_KEY")
+
+mix = {
+    "implementation": "Can you implement a matmul kernel in Triton", 
+    "conversion": "Convert this pytorch model to Triton", 
+    "explanation": "What memory access optimizations are being used here?"
+}
+
+ds = dataset({
+    "id": uuid,
+    "task": sample(mix),
+    "prompt": gen("write a prompt for {task}"),
+    "response": gen(f"write a response to {prompt}"),
+)}
+
+
+

chatan-0.1.0/pyproject.toml

@@ -0,0 +1,71 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "chatan"
+dynamic = ["version"]
+description = "Create synthetic datasets with LLM generators and samplers"
+readme = "README.md"
+requires-python = ">=3.8"
+license = "MIT"
+keywords = ["synthetic data", "llm", "dataset generation", "machine learning"]
+authors = [
+    { name = "Christian Reetz", email = "cdreetz@gmail.com" },
+]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+dependencies = [
+    "openai>=1.0.0",
+    "anthropic>=0.7.0",
+    "datasets>=2.0.0",
+    "pandas>=1.3.0",
+    "numpy>=1.20.0",
+    "pydantic>=2.0.0",
+]
+
+[dependency-groups]
+dev = [
+    "pytest>=7.0",
+    "black>=22.0",
+    "isort>=5.0",
+    "mypy>=1.0",
+    "pytest-cov>=4.0",
+]
+
+[project.urls]
+Documentation = "https://github.com/cdreetz/chatan#readme"
+Issues = "https://github.com/cdreetz/chatan/issues"
+Source = "https://github.com/cdreetz/chatan"
+
+[tool.hatch.version]
+path = "src/chatan/__init__.py"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/chatan"]
+
+[tool.black]
+line-length = 88
+target-version = ['py38']
+
+[tool.isort]
+profile = "black"
+multi_line_output = 3
+
+[tool.mypy]
+python_version = "3.8"
+warn_return_any = true
+warn_unused_configs = true
+disallow_untyped_defs = true

chatan-0.1.0/src/chatan/__init__.py

@@ -0,0 +1,9 @@
+"""Minos: Create synthetic datasets with LLM generators and samplers."""
+
+__version__ = "0.1.0"
+
+from .dataset import dataset
+from .generator import generator
+from .sampler import sample
+
+__all__ = ["dataset", "generator", "sample"]

chatan-0.1.0/src/chatan/dataset.py

@@ -0,0 +1,131 @@
+"""Dataset creation and manipulation."""
+
+from typing import Dict, Any, Union, Optional, List, Callable
+import pandas as pd
+from datasets import Dataset as HFDataset
+from .generator import GeneratorFunction
+from .sampler import SampleFunction
+
+
+class Dataset:
+    """Synthetic dataset generator."""
+    
+    def __init__(self, schema: Union[Dict[str, Any], str], n: int = 100):
+        """
+        Initialize dataset with schema.
+        
+        Args:
+            schema: Either a dict mapping column names to generators/samplers,
+                   or a string prompt for high-level dataset generation
+            n: Number of samples to generate
+        """
+        if isinstance(schema, str):
+            # High-level prompting - we'll implement this later
+            raise NotImplementedError("High-level prompting not yet implemented")
+        
+        self.schema = schema
+        self.n = n
+        self._data = None
+    
+    def generate(self, n: Optional[int] = None) -> pd.DataFrame:
+        """Generate the dataset."""
+        num_samples = n or self.n
+        
+        # Build dependency graph
+        dependencies = self._build_dependency_graph()
+        execution_order = self._topological_sort(dependencies)
+        
+        # Generate data
+        data = []
+        for i in range(num_samples):
+            row = {}
+            for column in execution_order:
+                value = self._generate_value(column, row)
+                row[column] = value
+            data.append(row)
+        
+        self._data = pd.DataFrame(data)
+        return self._data
+    
+    def _build_dependency_graph(self) -> Dict[str, List[str]]:
+        """Build dependency graph from schema."""
+        dependencies = {}
+        
+        for column, func in self.schema.items():
+            deps = []
+            if isinstance(func, GeneratorFunction):
+                # Extract column references from prompt template
+                import re
+                template = func.prompt_template
+                deps = re.findall(r'\{(\w+)\}', template)
+            
+            dependencies[column] = [dep for dep in deps if dep in self.schema]
+        
+        return dependencies
+    
+    def _topological_sort(self, dependencies: Dict[str, List[str]]) -> List[str]:
+        """Topologically sort columns by dependencies."""
+        visited = set()
+        temp_visited = set()
+        result = []
+        
+        def visit(column):
+            if column in temp_visited:
+                raise ValueError(f"Circular dependency detected involving {column}")
+            if column in visited:
+                return
+            
+            temp_visited.add(column)
+            for dep in dependencies.get(column, []):
+                visit(dep)
+            temp_visited.remove(column)
+            visited.add(column)
+            result.append(column)
+        
+        for column in self.schema:
+            visit(column)
+        
+        return result
+    
+    def _generate_value(self, column: str, context: Dict[str, Any]) -> Any:
+        """Generate a single value for a column."""
+        func = self.schema[column]
+        
+        if isinstance(func, (GeneratorFunction, SampleFunction)):
+            return func(context)
+        elif callable(func):
+            return func(context)
+        else:
+            # Static value
+            return func
+    
+    def to_pandas(self) -> pd.DataFrame:
+        """Convert to pandas DataFrame."""
+        if self._data is None:
+            self.generate()
+        return self._data
+    
+    def to_huggingface(self) -> HFDataset:
+        """Convert to HuggingFace Dataset."""
+        if self._data is None:
+            self.generate()
+        return HFDataset.from_pandas(self._data)
+    
+    def save(self, path: str, format: str = "parquet") -> None:
+        """Save dataset to file."""
+        if self._data is None:
+            self.generate()
+        
+        if format == "parquet":
+            self._data.to_parquet(path)
+        elif format == "csv":
+            self._data.to_csv(path, index=False)
+        elif format == "json":
+            self._data.to_json(path, orient="records")
+        else:
+            raise ValueError(f"Unsupported format: {format}")
+
+
+def dataset(schema: Union[Dict[str, Any], str], n: int = 100) -> Dataset:
+    """Create a synthetic dataset."""
+    return Dataset(schema, n)

chatan-0.1.0/src/chatan/generator.py

@@ -0,0 +1,94 @@
+"""LLM generators for synthetic data creation."""
+
+from typing import Dict, Any, Optional, Union, List
+import openai
+import anthropic
+from abc import ABC, abstractmethod
+
+
+class BaseGenerator(ABC):
+    """Base class for LLM generators."""
+    
+    @abstractmethod
+    def generate(self, prompt: str, **kwargs) -> str:
+        """Generate content from a prompt."""
+        pass
+
+
+class OpenAIGenerator(BaseGenerator):
+    """OpenAI GPT generator."""
+    
+    def __init__(self, api_key: str, model: str = "gpt-3.5-turbo", **kwargs):
+        self.client = openai.OpenAI(api_key=api_key)
+        self.model = model
+        self.default_kwargs = kwargs
+    
+    def generate(self, prompt: str, **kwargs) -> str:
+        """Generate content using OpenAI API."""
+        merged_kwargs = {**self.default_kwargs, **kwargs}
+        
+        response = self.client.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": prompt}],
+            **merged_kwargs
+        )
+        return response.choices[0].message.content.strip()
+
+
+class AnthropicGenerator(BaseGenerator):
+    """Anthropic Claude generator."""
+    
+    def __init__(self, api_key: str, model: str = "claude-3-sonnet-20240229", **kwargs):
+        self.client = anthropic.Anthropic(api_key=api_key)
+        self.model = model
+        self.default_kwargs = kwargs
+    
+    def generate(self, prompt: str, **kwargs) -> str:
+        """Generate content using Anthropic API."""
+        merged_kwargs = {**self.default_kwargs, **kwargs}
+        
+        response = self.client.messages.create(
+            model=self.model,
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=merged_kwargs.pop("max_tokens", 1000),
+            **merged_kwargs
+        )
+        return response.content[0].text.strip()
+
+
+class GeneratorFunction:
+    """Callable generator function for use in dataset schemas."""
+    
+    def __init__(self, generator: BaseGenerator, prompt_template: str):
+        self.generator = generator
+        self.prompt_template = prompt_template
+    
+    def __call__(self, context: Dict[str, Any]) -> str:
+        """Generate content with context substitution."""
+        prompt = self.prompt_template.format(**context)
+        result = self.generator.generate(prompt)
+        return result.strip() if isinstance(result, str) else result
+
+
+class GeneratorClient:
+    """Main interface for creating generators."""
+    
+    def __init__(self, provider: str, api_key: str, **kwargs):
+        if provider.lower() == "openai":
+            self._generator = OpenAIGenerator(api_key, **kwargs)
+        elif provider.lower() == "anthropic":
+            self._generator = AnthropicGenerator(api_key, **kwargs)
+        else:
+            raise ValueError(f"Unsupported provider: {provider}")
+    
+    def __call__(self, prompt_template: str) -> GeneratorFunction:
+        """Create a generator function."""
+        return GeneratorFunction(self._generator, prompt_template)
+
+
+# Factory function
+def generator(provider: str = "openai", api_key: Optional[str] = None, **kwargs) -> GeneratorClient:
+    """Create a generator client."""
+    if api_key is None:
+        raise ValueError("API key is required")
+    return GeneratorClient(provider, api_key, **kwargs)

chatan-0.1.0/src/chatan/sampler.py

@@ -0,0 +1,143 @@
+"""Sampling functions for synthetic data creation."""
+
+import random
+import uuid
+from datetime import datetime, timedelta
+from typing import Any, Dict, List, Union, Optional, Callable
+import pandas as pd
+from datasets import Dataset as HFDataset
+
+
+class SampleFunction:
+    """Base class for sampling functions."""
+    
+    def __call__(self, context: Dict[str, Any] = None) -> Any:
+        """Generate a sample value."""
+        raise NotImplementedError
+
+
+class ChoiceSampler(SampleFunction):
+    """Sample from a list of choices."""
+    
+    def __init__(self, choices: Union[List[Any], Dict[str, Any]]):
+        if isinstance(choices, dict):
+            self.choices = list(choices.keys())
+            self.weights = list(choices.values())
+        else:
+            self.choices = choices
+            self.weights = None
+    
+    def __call__(self, context: Dict[str, Any] = None) -> Any:
+        if self.weights:
+            return random.choices(self.choices, weights=self.weights, k=1)[0]
+        return random.choice(self.choices)
+
+
+class WeightedSampler(SampleFunction):
+    """Sample from weighted choices."""
+    
+    def __init__(self, choices: Dict[str, float]):
+        self.choices = list(choices.keys())
+        self.weights = list(choices.values())
+    
+    def __call__(self, context: Dict[str, Any] = None) -> Any:
+        return random.choices(self.choices, weights=self.weights, k=1)[0]
+
+
+class UUIDSampler(SampleFunction):
+    """Generate UUID strings."""
+    
+    def __call__(self, context: Dict[str, Any] = None) -> str:
+        return str(uuid.uuid4())
+
+
+class DatetimeSampler(SampleFunction):
+    """Sample random datetimes."""
+    
+    def __init__(self, start: str, end: str, format: str = "%Y-%m-%d"):
+        self.start = datetime.strptime(start, format)
+        self.end = datetime.strptime(end, format)
+        self.delta = self.end - self.start
+    
+    def __call__(self, context: Dict[str, Any] = None) -> datetime:
+        random_days = random.randint(0, self.delta.days)
+        return self.start + timedelta(days=random_days)
+
+
+class RangeSampler(SampleFunction):
+    """Sample from numeric ranges."""
+    
+    def __init__(self, start: Union[int, float], end: Union[int, float], 
+                 step: Optional[Union[int, float]] = None):
+        self.start = start
+        self.end = end
+        self.step = step
+        self.is_int = isinstance(start, int) and isinstance(end, int)
+    
+    def __call__(self, context: Dict[str, Any] = None) -> Union[int, float]:
+        if self.is_int:
+            return random.randint(self.start, self.end)
+        return random.uniform(self.start, self.end)
+
+
+class DatasetSampler(SampleFunction):
+    """Sample from existing dataset columns."""
+    
+    def __init__(self, dataset: Union[pd.DataFrame, HFDataset, Dict], 
+                 column: str, default: Optional[SampleFunction] = None):
+        if isinstance(dataset, pd.DataFrame):
+            self.values = dataset[column].tolist()
+        elif isinstance(dataset, HFDataset):
+            self.values = dataset[column]
+        elif isinstance(dataset, dict):
+            self.values = dataset[column]
+        else:
+            raise ValueError("Unsupported dataset type")
+        
+        self.default = default
+    
+    def __call__(self, context: Dict[str, Any] = None) -> Any:
+        if not self.values and self.default:
+            return self.default(context)
+        return random.choice(self.values)
+
+
+# Factory functions for the sample namespace
+class SampleNamespace:
+    """Namespace for sampling functions."""
+    
+    @staticmethod
+    def choice(choices: Union[List[Any], Dict[str, Any]]) -> ChoiceSampler:
+        """Sample from choices."""
+        return ChoiceSampler(choices)
+    
+    @staticmethod
+    def weighted(choices: Dict[str, float]) -> WeightedSampler:
+        """Sample from weighted choices."""
+        return WeightedSampler(choices)
+    
+    @staticmethod
+    def uuid() -> UUIDSampler:
+        """Generate UUIDs."""
+        return UUIDSampler()
+    
+    @staticmethod
+    def datetime(start: str, end: str, format: str = "%Y-%m-%d") -> DatetimeSampler:
+        """Sample random datetimes."""
+        return DatetimeSampler(start, end, format)
+    
+    @staticmethod
+    def range(start: Union[int, float], end: Union[int, float], 
+              step: Optional[Union[int, float]] = None) -> RangeSampler:
+        """Sample from numeric ranges."""
+        return RangeSampler(start, end, step)
+    
+    @staticmethod
+    def from_dataset(dataset: Union[pd.DataFrame, HFDataset, Dict], 
+                     column: str, default: Optional[SampleFunction] = None) -> DatasetSampler:
+        """Sample from existing dataset."""
+        return DatasetSampler(dataset, column, default)
+
+
+# Export the sample namespace
+sample = SampleNamespace()