ASTChunk

https://github.com/yilinjz/astchunk
Admin
ASTChunk is an AST-based code chunking tool that intelligently divides source code into meaningful...
Tokens:28,937
Snippets:167
Trust Score:5.5
License:MIT
Update:6 months ago
Show doc for...
Tokens:
Raw
### Quick Start with ASTChunk

Source: https://github.com/yilinjz/astchunk/blob/main/README.md

Basic usage example for initializing the builder and chunking a Python code string.

```python
from astchunk import ASTChunkBuilder

# Your source code
code = """
def fibonacci(n):
    if n <= 1:
        return n
    return fibonacci(n-1) + fibonacci(n-2)

class Calculator:
    def add(self, a, b):
        return a + b
    
    def multiply(self, a, b):
        return a * b
"""

# Initialize the chunk builder
configs = {
    "max_chunk_size": 100,             # Maximum non-whitespace characters per chunk
    "language": "python",              # Supported: python, java, csharp, typescript
    "metadata_template": "default"     # Metadata format for output
}
chunk_builder = ASTChunkBuilder(**configs)

# Create chunks
chunks = chunk_builder.chunkify(code)

# Each chunk contains content and metadata
for i, chunk in enumerate(chunks):
    print(f"[Chunk {i+1}]")
    print(f"{chunk['content']}")
    print(f"Metadata: {chunk['metadata']}")
    print("-" * 50)
```

--------------------------------

### Install ASTChunk

Source: https://github.com/yilinjz/astchunk/blob/main/README.md

Installation commands for the library via pip or source.

```bash
pip install astchunk
```

```bash
git clone git@github.com:yilinjz/astchunk.git
pip install -e .
```

--------------------------------

### Checkpointing Setup

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/fixed_chunking_results.txt

Configures checkpointing path, frequency, and maximum number of checkpoints to keep. Determines if checkpointing is possible based on distributed environment and filesystem.

```python
assert not (exists(checkpoint_path) ^ exists(checkpoint_every))
        self.checkpoint_path = checkpoint_path
        self.checkpoint_every = checkpoint_every
        self.max_checkpoints_keep = max_checkpoints_keep

        self.can_checkpoint = self.is_local_main if isinstance(checkpoint_fs, LocalFileSystem) else self.is_main

        if exists(checkpoint_path) and self.can_checkpoint:
            bucket = url_to_bucket(checkpoint_path)

            if not self.fs.exists(bucket):
                self.fs.mkdir(bucket)

            self.load_from_checkpoint_folder()
```

--------------------------------

### Create Single-Use Configs for chunkify()

Source: https://github.com/yilinjz/astchunk/blob/main/README.md

Demonstrates how to create and use single-use configurations for the chunkify() method, allowing for specific arguments per call. The example shows saving the generated chunks to separate files.

```python
single_use_configs = {
    "repo_level_metadata": {
        "filepath": "example.py"
    },
    "chunk_expansion": True
}

chunks = chunk_builder.chunkify(code, **single_use_configs)

# Save chunks to separate files
for i, chunk in enumerate(chunks):
    with open(f"chunk_{i+1}.py", "w") as f:
        f.write(chunk['content'])
```

--------------------------------

### Install ASTChunk Dependencies

Source: https://github.com/yilinjz/astchunk/blob/main/README.md

Required dependencies for AST parsing, which are typically handled automatically.

```bash
# Core dependencies (automatically installed)
pip install numpy pyrsistent tree-sitter
pip install tree-sitter-python tree-sitter-java tree-sitter-c-sharp tree-sitter-typescript
```

--------------------------------

### Setup Cosine Annealing and Warmup Schedulers

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_results.txt

Initializes learning rate schedulers, including CosineAnnealingLR and LinearWarmup, based on specified parameters. Handles cases where only one scheduler is needed.

```python
scheduler = warmup_scheduler = None

            if exists(unet_cosine_decay_max_steps):
                scheduler = CosineAnnealingLR(optimizer, T_max = unet_cosine_decay_max_steps)

            if exists(unet_warmup_steps):
                warmup_scheduler = warmup.LinearWarmup(optimizer, warmup_period = unet_warmup_steps)

                if not exists(scheduler):
                    scheduler = LambdaLR(optimizer, lr_lambda = lambda step: 1.0)
```

--------------------------------

### Process Files with ASTChunk

Source: https://github.com/yilinjz/astchunk/blob/main/README.md

Example of reading a file from disk to process with the chunker.

```python
# Process a single file
with open("example.py", "r") as f:
    code = f.read()
```

--------------------------------

### Initialize Trainer State and Device Placement

Source: https://github.com/yilinjz/astchunk/blob/main/examples/source_code.txt

Initializes trainer state, registers step buffers, and moves the model and trainer to the appropriate device. Handles checkpointing setup based on provided paths and filesystem.

```python
# gradient clipping if needed

        self.max_grad_norm = max_grad_norm

        # step tracker and misc

        self.register_buffer('steps', torch.tensor([0] * self.num_unets))

        self.verbose = verbose

        # automatic set devices based on what accelerator decided

        self.imagen.to(self.device)
        self.to(self.device)

        # checkpointing

        assert not (exists(checkpoint_path) ^ exists(checkpoint_every))
        self.checkpoint_path = checkpoint_path
        self.checkpoint_every = checkpoint_every
        self.max_checkpoints_keep = max_checkpoints_keep

        self.can_checkpoint = self.is_local_main if isinstance(checkpoint_fs, LocalFileSystem) else self.is_main

        if exists(checkpoint_path) and self.can_checkpoint:
            bucket = url_to_bucket(checkpoint_path)

            if not self.fs.exists(bucket):
                self.fs.mkdir(bucket)

            self.load_from_checkpoint_folder()

        # only allowing training for unet

        self.only_train_unet_number = only_train_unet_number
        self.prepared = False
```

--------------------------------

### Direct AST Processing Example

Source: https://github.com/yilinjz/astchunk/blob/main/README.md

Provides a placeholder for direct AST processing using ASTNode and ASTChunk. Refer to the API documentation for detailed usage instructions on custom processing with these components.

```python
from astchunk.astnode import ASTNode
from astchunk.astchunk import ASTChunk

# Work directly with AST nodes and chunks for custom processing
# (See API documentation for detailed usage)
```

--------------------------------

### ImagenTrainer Initialization

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_with_expansion_results.txt

Initializes the ImagenTrainer class, setting up parameters for training an Imagen model. It handles distributed training setup and accelerator configuration.

```python
class ImagenTrainer(nn.Module):
    locked = False

    def __init__(
        self,
        imagen = None,
        imagen_checkpoint_path = None,
        use_ema = True,
        lr = 1e-4,
        eps = 1e-8,
        beta1 = 0.9,
        beta2 = 0.99,
        max_grad_norm = None,
        group_wd_params = True,
        warmup_steps = None,
        cosine_decay_max_steps = None,
        only_train_unet_number = None,
        fp16 = False,
        precision = None,
        split_batches = True,
        dl_tuple_output_keywords_names = ('images', 'text_embeds', 'text_masks', 'cond_images'),
        verbose = True,
        split_valid_fraction = 0.025,
        split_valid_from_train = False,
        split_random_seed = 42,
        checkpoint_path = './',
        checkpoint_every = None,
        checkpoint_fs = None,
        fs_kwargs: dict = None,
        max_checkpoints_keep = 20,
        **kwargs
    ):
        super().__init__()
        assert not ImagenTrainer.locked, 'ImagenTrainer can only be initialized once per process - for the sake of distributed training, you will now have to create a separate script to train each unet (or a script that accepts unet number as an argument)'
        assert exists(imagen) ^ exists(imagen_checkpoint_path), 'either imagen instance is passed into the trainer, or a checkpoint path that contains the imagen config'

        # determine filesystem, using fsspec, for saving to local filesystem or cloud

        self.fs = checkpoint_fs

        if not exists(self.fs):
            fs_kwargs = default(fs_kwargs, {})
            self.fs, _ = url_to_fs(default(checkpoint_path, './'), **fs_kwargs)

        assert isinstance(imagen, (Imagen, ElucidatedImagen))
        ema_kwargs, kwargs = groupby_prefix_and_trim('ema_', kwargs)

        # elucidated or not

        self.is_elucidated = isinstance(imagen, ElucidatedImagen)

        # create accelerator instance

        accelerate_kwargs, kwargs = groupby_prefix_and_trim('accelerate_', kwargs)

        assert not (fp16 and exists(precision)), 'either set fp16 = True or forward the precision ("fp16", "bf16") to Accelerator'
        accelerator_mixed_precision = default(precision, 'fp16' if fp16 else 'no')

        self.accelerator = Accelerator(**{
            'split_batches': split_batches,
            'mixed_precision': accelerator_mixed_precision,
            'kwargs_handlers': [DistributedDataParallelKwargs(find_unused_parameters = True)]
        , **accelerate_kwargs})

        ImagenTrainer.locked = self.is_distributed

        # cast data to fp16 at training time if needed

        self.cast_half_at_training = accelerator_mixed_precision == 'fp16'

        # grad scaler must be managed outside of accelerator

        grad_scaler_enabled = fp16

        # imagen, unets and ema unets

        self.imagen = imagen
        self.num_unets = len(self.imagen.unets)

        self.use_ema = use_ema and self.is_main
```

--------------------------------

### Check if a string begins with a prefix

Source: https://github.com/yilinjz/astchunk/blob/main/examples/source_code.txt

A simple helper function to determine if a string starts with a specified prefix.

```python
def string_begins_with(prefix, str):
    return str.startswith(prefix)
```

--------------------------------

### Get Training Steps

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/fixed_chunking_results.txt

Retrieves the number of training steps taken for a specific UNet.

```python
    def num_steps_taken(self, unet_number = None):
        if self.num_unets == 1:
            unet_number = default(unet_number, 1)

        return self.steps[unet_number - 1].item()
```

--------------------------------

### Group dictionary by key prefix

Source: https://github.com/yilinjz/astchunk/blob/main/examples/source_code.txt

Uses `group_dict_by_key` to partition a dictionary into two: one with keys starting with a given prefix, and another with the rest.

```python
def group_by_key_prefix(prefix, d):
    return group_dict_by_key(partial(string_begins_with, prefix), d)
```

--------------------------------

### Get Accelerator Device Property

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_with_expansion_results.txt

Provides access to the device (e.g., 'cuda', 'cpu') managed by the accelerator. Useful for ensuring tensors and models are on the correct device.

```python
@property
    def device(self):
        return self.accelerator.device
```

--------------------------------

### Initialize ASTChunkBuilder with Custom Metadata Templates

Source: https://github.com/yilinjz/astchunk/blob/main/README.md

Illustrates initializing the ASTChunkBuilder with specific metadata templates for use cases like repoeval and swebench-lite. This allows for tailored metadata generation during code chunking.

```python
# For repoeval
repoeval_builder = ASTChunkBuilder(
    max_chunk_size=2000,
    language="python",
    metadata_template="coderagbench-repoeval"
)

# For swebench-lite
swebench_builder = ASTChunkBuilder(
    max_chunk_size=2000,
    language="python",
    metadata_template="coderagbench-swebench-lite"
)
```

--------------------------------

### Get Learning Rate for a Specific UNet

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_with_expansion_results.txt

Retrieves the current learning rate for a specified UNet. It validates the UNet number and accesses the corresponding optimizer to get the LR.

```python
def get_lr(self, unet_number):
        self.validate_unet_number(unet_number)
        unet_index = unet_number - 1
```

--------------------------------

### Initialize ASTChunkBuilder for SWE-bench

Source: https://context7.com/yilinjz/astchunk/llms.txt

Configures the builder with specific metadata templates for SWE-bench Lite tasks.

```python
swebench_builder = ASTChunkBuilder(
    max_chunk_size=2000,
    language="python",
    metadata_template="coderagbench-swebench-lite"
)

chunks = swebench_builder.chunkify(
    code,
    repo_level_metadata={
        "instance_id": "issue-123",
        "filename": "example.py"
    }
)
```

--------------------------------

### Get Current Learning Rate

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_results.txt

Retrieves the current learning rate for a specified UNet model from its optimizer.

```python
self.validate_unet_number(unet_number)
        unet_index = unet_number - 1

        optim = getattr(self, f'optim{unet_index}')

        return optim.param_groups[0]['lr']
```

--------------------------------

### Execute Training Step

Source: https://github.com/yilinjz/astchunk/blob/main/examples/source_code.txt

Prepares the trainer and initializes the training iterator for a single training step.

```python
def train_step(self, *, unet_number = None, **kwargs):
        if not self.prepared:
            self.prepare()
        self.create_train_iter()

        kwargs = {'unet_number': unet_number, **kwargs}
```

--------------------------------

### Configure metadata templates

Source: https://context7.com/yilinjz/astchunk/llms.txt

Apply different metadata templates to suit specific integration requirements.

```python
from astchunk import ASTChunkBuilder

# Default template - general purpose
default_builder = ASTChunkBuilder(
    max_chunk_size=1500,
    language="python",
    metadata_template="default"
)

code = "def example(): pass"

# Default metadata includes:
# - filepath, chunk_size, line_count, start_line_no, end_line_no, node_count
chunks = default_builder.chunkify(
    code,
    repo_level_metadata={"filepath": "src/example.py"}
)
print(chunks[0]['metadata'])
# {'filepath': 'src/example.py', 'chunk_size': 18, 'line_count': 1,
#  'start_line_no': 0, 'end_line_no': 0, 'node_count': 1}

# CodeRAGBench RepoEval template
repoeval_builder = ASTChunkBuilder(
    max_chunk_size=2000,
    language="python",
    metadata_template="coderagbench-repoeval"
)

chunks = repoeval_builder.chunkify(
    code,
    repo_level_metadata={
        "fpath_tuple": ["myrepo", "src", "example.py"],
        "repo": "myrepo"
    }
)
```

--------------------------------

### Process Codebase with ASTChunk

Source: https://context7.com/yilinjz/astchunk/llms.txt

Demonstrates a full pipeline for reading files, chunking with expansion and overlap, and writing results to a file.

```python
from astchunk import ASTChunkBuilder

def process_codebase(source_files, output_file, max_chunk_size=1800):
    """Process multiple source files into AST chunks."""

    chunk_builder = ASTChunkBuilder(
        max_chunk_size=max_chunk_size,
        language="python",
        metadata_template="default"
    )

    all_chunks = []

    for filepath in source_files:
        with open(filepath, 'r', encoding='utf-8') as f:
            code = f.read()

        chunks = chunk_builder.chunkify(
            code,
            repo_level_metadata={"filepath": filepath},
            chunk_expansion=True,
            chunk_overlap=1
        )

        all_chunks.extend(chunks)
        print(f"Processed {filepath}: {len(chunks)} chunks")

    # Write results
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(f"Total chunks: {len(all_chunks)}\n")
        f.write("=" * 80 + "\n\n")

        for i, chunk in enumerate(all_chunks, 1):
            content = chunk['content']
            metadata = chunk['metadata']

            header = f"--- Chunk {i} ({metadata['chunk_size']} chars) ---\n"
            f.write(header)
            f.write(f"File: {metadata['filepath']}\n")
            f.write(f"Lines: {metadata['start_line_no']}-{metadata['end_line_no']}\n\n")
            f.write(content)
            f.write("\n" + "-" * 40 + "\n\n")

    return all_chunks

# Usage
# chunks = process_codebase(
#     source_files=['src/main.py', 'src/utils.py'],
#     output_file='chunks_output.txt'
# )
```

--------------------------------

### Get Number of Training Steps

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_with_expansion_results.txt

Retrieves the number of training steps taken for a specific UNet. If only one UNet exists, it defaults to UNet 1.

```python
def num_steps_taken(self, unet_number = None):
        if self.num_unets == 1:
            unet_number = default(unet_number, 1)

        return self.steps[unet_number - 1].item()
```

--------------------------------

### Initialize ImagenTrainer with Checkpointing and Device Settings

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_with_expansion_results.txt

Configures checkpointing paths, frequency, and retention policies. It also handles loading from existing checkpoints and setting the model to the appropriate device based on the accelerator.

```python
self.checkpoint_path = checkpoint_path
        self.checkpoint_every = checkpoint_every
        self.max_checkpoints_keep = max_checkpoints_keep

        self.can_checkpoint = self.is_local_main if isinstance(checkpoint_fs, LocalFileSystem) else self.is_main

        if exists(checkpoint_path) and self.can_checkpoint:
            bucket = url_to_bucket(checkpoint_path)

            if not self.fs.exists(bucket):
                self.fs.mkdir(bucket)

            self.load_from_checkpoint_folder()

        # only allowing training for unet

        self.only_train_unet_number = only_train_unet_number
        self.prepared = False


    def prepare(self):
        assert not self.prepared, f'The trainer is allready prepared'
        self.validate_and_set_unet_being_trained(self.only_train_unet_number)
        self.prepared = True
```

--------------------------------

### Initialize ASTChunkBuilder and Chunkify Code

Source: https://context7.com/yilinjz/astchunk/llms.txt

Initializes the ASTChunkBuilder with configuration options and processes source code into structured chunks. Ensure the 'language' in configs matches the source code.

```python
from astchunk import ASTChunkBuilder

# Initialize the chunk builder with configuration
configs = {
    "max_chunk_size": 1800,     # Maximum non-whitespace characters per chunk
    "language": "python",       # Supported: python, java, csharp, typescript
    "metadata_template": "default"  # Metadata format for output
}
chunk_builder = ASTChunkBuilder(**configs)

# Source code to chunk
code = """
import numpy as np

class Calculator:
    def __init__(self, precision=2):
        self.precision = precision

    def add(self, a, b):
        return round(a + b, self.precision)

    def multiply(self, a, b):
        return round(a * b, self.precision)

def fibonacci(n):
    if n <= 1:
        return n
    return fibonacci(n-1) + fibonacci(n-2)
"""

# Create chunks
chunks = chunk_builder.chunkify(code)

# Process results
for i, chunk in enumerate(chunks, 1):
    content = chunk['content']
    metadata = chunk['metadata']
    print(f"Chunk {i}: {metadata['chunk_size']} chars, lines {metadata['start_line_no']}-{metadata['end_line_no']}")
    print(content)
    print("-" * 50)

# Output:
# Chunk 1: 287 chars, lines 0-17
# import numpy as np
#
# class Calculator:
#     def __init__(self, precision=2):
#         ...

```

--------------------------------

### Get Sorted Checkpoints

Source: https://github.com/yilinjz/astchunk/blob/main/examples/source_code.txt

Retrieves a list of all checkpoint files in the specified directory, sorted in descending order by their step number. This is useful for loading the latest checkpoint.

```python
glob_pattern = os.path.join(self.checkpoint_path, '*.pt')
checkpoints = self.fs.glob(glob_pattern)
sorted_checkpoints = sorted(checkpoints, key = lambda x: int(str(x).split('.')[-2]), reverse = True)
return sorted_checkpoints
```

--------------------------------

### Check if Distributed Training is Enabled

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_with_expansion_results.txt

Determines if the current training setup is distributed across multiple processes or GPUs. Returns true if distributed training is active.

```python
@property
    def is_distributed(self):
        return not (self.accelerator.distributed_type == DistributedType.NO and self.accelerator.num_processes == 1)
```

--------------------------------

### Advanced Chunking Configuration

Source: https://github.com/yilinjz/astchunk/blob/main/README.md

Demonstrates how to apply repository-level metadata, overlapping, and chunk expansion.

```python

# Add repo-level metadata
configs['repo_level_metadata'] = {
    "filepath": "src/calculator.py"
}

# Enable overlapping between chunks
configs['chunk_overlap'] = 1

# Add chunk expansion (metadata headers)
configs['chunk_expansion'] = True

# NOTE: max_chunk_size apply to the chunks before overlapping or chunk expansion.
# The final chunk size after overlapping or chunk expansion may exceed max_chunk_size.


# Extend current code for illustration
code += """
def divide(self, a, b):
    if b == 0:
        raise ValueError("Cannot divide by zero")
    return a / b

# This is a comment
# Another comment

def subtract(self, a, b):
    return a - b

def exponent(self, a, b):
    return a ** b
"""


# Create chunks
chunks = chunk_builder.chunkify(code, **configs)

for i, chunk in enumerate(chunks):
    print(f"[Chunk {i+1}]")
    print(f"{chunk['content']}")
    print(f"Metadata: {chunk['metadata']}")
    print("-" * 50)
```

--------------------------------

### POST /sample

Source: https://github.com/yilinjz/astchunk/blob/main/examples/source_code.txt

Generates samples from the model using chunked processing.

```APIDOC
## POST /sample

### Description
Generates samples from the underlying Imagen model. This method is decorated with @torch.no_grad and supports chunked sampling.

### Method
POST

### Endpoint
/sample

### Parameters
#### Request Body
- **use_non_ema** (bool) - Optional - If true, uses non-EMA weights for sampling.
- **use_tqdm** (bool) - Optional - Controls progress bar display.
```

--------------------------------

### Get Unwrapped UNet Model

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_with_expansion_results.txt

Retrieves the UNet model being trained, unwrapped from any distributed or mixed-precision wrappers applied by the accelerator. This provides access to the original model parameters.

```python
@property
    def unwrapped_unet(self):
        return self.accelerator.unwrap_model(self.unet_being_trained)
```

--------------------------------

### Sample from model

Source: https://github.com/yilinjz/astchunk/blob/main/examples/source_code.txt

Generates samples using the internal imagen model, supporting chunked sampling and EMA toggling.

```python
@torch.no_grad()
@cast_torch_tensor
@imagen_sample_in_chunks
def sample(self, *args, **kwargs):
    context = nullcontext if  kwargs.pop('use_non_ema', False) else self.use_ema_unets

    self.print_untrained_unets()

    if not self.is_main:
        kwargs['use_tqdm'] = False

    with context():
        output = self.imagen.sample(*args, device = self.device, **kwargs)

    return output
```

--------------------------------

### Get Current Learning Rate for a UNet

Source: https://github.com/yilinjz/astchunk/blob/main/examples/source_code.txt

Retrieves the current learning rate for a specified UNet. It validates the UNet number and accesses the optimizer's parameter group.

```python
def get_lr(self, unet_number):
        self.validate_unet_number(unet_number)
        unet_index = unet_number - 1

        optim = getattr(self, f'optim{unet_index}')

        return optim.param_groups[0]['lr']
```

--------------------------------

### Wrap UNet with Accelerator

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/fixed_chunking_results.txt

Prepares a specific UNet and its associated optimizer/scheduler using the accelerator.

```python
    def wrap_unet(self, unet_number):
        if hasattr(self, 'one_unet_wrapped'):
            return

        unet = self.imagen.get_unet(unet_number)
        unet_index = unet_number - 1

        optimizer = getattr(self, f'optim{unet_index}')
        scheduler = getattr(self, f'scheduler{unet_index}')

        if self.train_dl:
            self.unet_being_trained, self.train_dl, optimizer = self.accelerator.prepare(unet, self.train_dl, optimizer)
        else:
            self.unet_being_trained, optimizer = self.accelerator.prepare(unet, optimizer)

        if exists(scheduler):
            scheduler = self.accelerator.prepare(scheduler)

        setattr(self, f'optim{unet_index}', optimizer)
        setattr(self, f'scheduler{unet_index}', scheduler)

        self.one_unet_wrapped = True
```

--------------------------------

### Check if Process is Local Main

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_with_expansion_results.txt

Returns true if the current process is the main process on the local machine, even in a distributed setup. Useful for actions that should only happen once per machine.

```python
@property
    def is_local_main(self):
        return self.accelerator.is_local_main_process
```

--------------------------------

### Forward Pass Initialization

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/fixed_chunking_results.txt

Prepares the model for a forward pass by validating the U-Net number and setting up the accelerator and scaler.

```python
@partial(cast_torch_tensor, cast_fp16 = True)
    def forward(
        self,
        *args,
        unet_number = None,
        max_batch_size = None,
        **kwargs
    ):
        unet_number = self.validate_unet_number(unet_number)
        self.validate_and_set_unet_being_trained(unet_number)
        self.set_accelerator_scaler(unet_number)

        assert not exists(self.only_train_unet_number) or self.only_train_unet_number == unet_number, f'you can only train unet #{self.only_train_unet_number}'

        total_loss = 0.
```

--------------------------------

### Initialize ASTChunkBuilder for different languages

Source: https://context7.com/yilinjz/astchunk/llms.txt

Configure builders for specific programming languages with defined maximum chunk sizes.

```python
# Java builder with larger chunks
java_builder = ASTChunkBuilder(
    max_chunk_size=2000,
    language="java",
    metadata_template="default"
)

# TypeScript/TSX builder
ts_builder = ASTChunkBuilder(
    max_chunk_size=1800,
    language="typescript",
    metadata_template="default"
)

# C# builder
csharp_builder = ASTChunkBuilder(
    max_chunk_size=2000,
    language="csharp",
    metadata_template="default"
)
```

--------------------------------

### Initialize Trainer Preparation State

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_results.txt

Sets the initial state for trainer preparation, ensuring it's not already prepared.

```python
self.prepared = False
```

--------------------------------

### Wrap and Prepare a Specific UNet for Training

Source: https://github.com/yilinjz/astchunk/blob/main/examples/source_code.txt

Wraps a specific UNet for training, preparing it along with the training data loader and optimizer using the accelerator. Avoids re-wrapping if already done.

```python
def wrap_unet(self, unet_number):
        if hasattr(self, 'one_unet_wrapped'):
            return

        unet = self.imagen.get_unet(unet_number)
        unet_index = unet_number - 1

        optimizer = getattr(self, f'optim{unet_index}')
        scheduler = getattr(self, f'scheduler{unet_index}')

        if self.train_dl:
            self.unet_being_trained, self.train_dl, optimizer = self.accelerator.prepare(unet, self.train_dl, optimizer)
        else:
            
```

--------------------------------

### Move Model to Device

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_results.txt

Ensures that the main model and the trainer itself are placed on the correct computation device.

```python
self.imagen.to(self.device)
        self.to(self.device)
```

--------------------------------

### ImagenTrainer Initialization

Source: https://github.com/yilinjz/astchunk/blob/main/examples/source_code.txt

Initializes the trainer with support for EMA, distributed training, and mixed precision.

```python
class ImagenTrainer(nn.Module):
    locked = False

    def __init__(
        self,
        imagen = None,
        imagen_checkpoint_path = None,
        use_ema = True,
        lr = 1e-4,
        eps = 1e-8,
        beta1 = 0.9,
        beta2 = 0.99,
        max_grad_norm = None,
        group_wd_params = True,
        warmup_steps = None,
        cosine_decay_max_steps = None,
        only_train_unet_number = None,
        fp16 = False,
        precision = None,
        split_batches = True,
        dl_tuple_output_keywords_names = ('images', 'text_embeds', 'text_masks', 'cond_images'),
        verbose = True,
        split_valid_fraction = 0.025,
        split_valid_from_train = False,
        split_random_seed = 42,
        checkpoint_path = None,
        checkpoint_every = None,
        checkpoint_fs = None,
        fs_kwargs: dict = None,
        max_checkpoints_keep = 20,
        **kwargs
    ):
        super().__init__()
        assert not ImagenTrainer.locked, 'ImagenTrainer can only be initialized once per process - for the sake of distributed training, you will now have to create a separate script to train each unet (or a script that accepts unet number as an argument)'
        assert exists(imagen) ^ exists(imagen_checkpoint_path), 'either imagen instance is passed into the trainer, or a checkpoint path that contains the imagen config'

        # determine filesystem, using fsspec, for saving to local filesystem or cloud

        self.fs = checkpoint_fs

        if not exists(self.fs):
            fs_kwargs = default(fs_kwargs, {})
            self.fs, _ = url_to_fs(default(checkpoint_path, './'), **fs_kwargs)

        assert isinstance(imagen, (Imagen, ElucidatedImagen))
        ema_kwargs, kwargs = groupby_prefix_and_trim('ema_', kwargs)

        # elucidated or not

        self.is_elucidated = isinstance(imagen, ElucidatedImagen)

        # create accelerator instance

        accelerate_kwargs, kwargs = groupby_prefix_and_trim('accelerate_', kwargs)

        assert not (fp16 and exists(precision)), 'either set fp16 = True or forward the precision ("fp16", "bf16") to Accelerator'
        accelerator_mixed_precision = default(precision, 'fp16' if fp16 else 'no')

        self.accelerator = Accelerator(**{
            'split_batches': split_batches,
            'mixed_precision': accelerator_mixed_precision,
            'kwargs_handlers': [DistributedDataParallelKwargs(find_unused_parameters = True)]
        , **accelerate_kwargs})

        ImagenTrainer.locked = self.is_distributed

        # cast data to fp16 at training time if needed

        self.cast_half_at_training = accelerator_mixed_precision == 'fp16'

        # grad scaler must be managed outside of accelerator

        grad_scaler_enabled = fp16

        # imagen, unets and ema unets

        self.imagen = imagen
        self.num_unets = len(self.imagen.unets)

        self.use_ema = use_ema and self.is_main
        self.ema_unets = nn.ModuleList([])

        # keep track of what unet is being trained on
        # only going to allow 1 unet training at a time

        self.ema_unet_being_trained_index = -1 # keeps track of which ema unet is being trained on

        # data related functions

        self.train_dl_iter = None
        self.train_dl = None

        self.valid_dl_iter = None
        self.valid_dl = None

        self.dl_tuple_output_keywords_names = dl_tuple_output_keywords_names

        # auto splitting validation from training, if dataset is passed in

        self.split_valid_from_train = split_valid_from_train

        assert 0 <= split_valid_fraction <= 1, 'split valid fraction must be between 0 and 1'
```

--------------------------------

### Add Training Dataset and Create Dataloader

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_results.txt

Adds a training dataset and optionally splits it to create a validation dataloader. Configures batch size and other dataloader arguments. Asserts that the training dataloader has not already been added.

```python
def add_train_dataset(self, ds = None, *, batch_size, **dl_kwargs):
        if not exists(ds):
            return

        assert not exists(self.train_dl), 'training dataloader was already added'

        valid_ds = None
        if self.split_valid_from_train:
            train_size = int((1 - self.split_valid_fraction) * len(ds))
            valid_size = len(ds) - train_size

            ds, valid_ds = random_split(ds, [train_size, valid_size], generator = torch.Generator().manual_seed(self.split_random_seed))
            self.print(f'training with dataset of {len(ds)} samples and validating with randomly splitted {len(valid_ds)} samples')
```

--------------------------------

### Perform Training Step

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_results.txt

Executes a single training step, including preparing the trainer, creating the training iterator, and updating the model. Returns the calculated loss.

```python
if not self.prepared:
    self.prepare()
self.create_train_iter()

kwargs = {'unet_number': unet_number, **kwargs}
loss = self.step_with_dl_iter(self.train_dl_iter, **kwargs)
self.update(unet_number = unet_number)
return loss
```

--------------------------------

### Wrap UNet and Prepare Optimizers/Schedulers

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_with_expansion_results.txt

Wraps a specific UNet model and prepares its optimizer and scheduler using the accelerator. This method ensures that only one UNet is wrapped and prepared at a time to avoid conflicts.

```python
def wrap_unet(self, unet_number):
        if hasattr(self, 'one_unet_wrapped'):
            return

        unet = self.imagen.get_unet(unet_number)
        unet_index = unet_number - 1

        optimizer = getattr(self, f'optim{unet_index}')
        scheduler = getattr(self, f'scheduler{unet_index}')

        if self.train_dl:
            self.unet_being_trained, self.train_dl, optimizer = self.accelerator.prepare(unet, self.train_dl, optimizer)
        else:
            self.unet_being_trained, optimizer = self.accelerator.prepare(unet, optimizer)

        if exists(scheduler):
            scheduler = self.accelerator.prepare(scheduler)

        setattr(self, f'optim{unet_index}', optimizer)
        setattr(self, f'scheduler{unet_index}', scheduler)

        self.one_unet_wrapped = True
```

--------------------------------

### Register Training and Validation Datasets

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_with_expansion_results.txt

Configures training and validation data loaders, optionally splitting the dataset if configured.

```python
        assert not exists(self.train_dl), 'training dataloader was already added'

        valid_ds = None
        if self.split_valid_from_train:
            train_size = int((1 - self.split_valid_fraction) * len(ds))
            valid_size = len(ds) - train_size

            ds, valid_ds = random_split(ds, [train_size, valid_size], generator = torch.Generator().manual_seed(self.split_random_seed))
            self.print(f'training with dataset of {len(ds)} samples and validating with randomly splitted {len(valid_ds)} samples')

        dl = DataLoader(ds, batch_size = batch_size, **dl_kwargs)
        self.add_train_dataloader(dl)

        if not self.split_valid_from_train:
            return

        self.add_valid_dataset(valid_ds, batch_size = batch_size, **dl_kwargs)

    def add_valid_dataset(self, ds, *, batch_size, **dl_kwargs):
        if not exists(ds):
            return

        assert not exists(self.valid_dl), 'validation dataloader was already added'

        dl = DataLoader(ds, batch_size = batch_size, **dl_kwargs)
        self.add_valid_dataloader(dl)
```

--------------------------------

### Prepare Trainer and Validate UNet

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_results.txt

Prepares the trainer by validating and setting the UNet model to be trained, preventing re-preparation.

```python
assert not self.prepared, f'The trainer is allready prepared'
        self.validate_and_set_unet_being_trained(self.only_train_unet_number)
        self.prepared = True
```

--------------------------------

### Prepare Trainer for Training

Source: https://github.com/yilinjz/astchunk/blob/main/examples/source_code.txt

Marks the trainer as prepared after validating and setting the UNet to be trained. Asserts that the trainer has not been prepared before.

```python
def prepare(self):
        assert not self.prepared, f'The trainer is allready prepared'
        self.validate_and_set_unet_being_trained(self.only_train_unet_number)
        self.prepared = True
```

--------------------------------

### Preprocess Code for Size Calculation

Source: https://github.com/yilinjz/astchunk/blob/main/README.md

Demonstrates preprocessing code to efficiently calculate its size in non-whitespace characters. This involves encoding the code to bytes and using `preprocess_nws_count` to create a cumulative sum array for quick lookups.

```python
from astchunk.preprocessing import preprocess_nws_count, get_nws_count, ByteRange

# Preprocess code for efficient size calculation
code_bytes = code.encode('utf-8')
nws_cumsum = preprocess_nws_count(code_bytes)

# Get non-whitespace character count for any byte range
byte_range = ByteRange(0, 100)  # First 100 bytes
char_count = get_nws_count(nws_cumsum, byte_range)
```

--------------------------------

### Configure Optimizers and Schedulers for UNets

Source: https://github.com/yilinjz/astchunk/blob/main/examples/source_code.txt

Sets up individual optimizers, EMA, gradient scalers, and learning rate schedulers for each UNet. Supports cosine annealing and linear warmup. Requires `Adam`, `EMA`, `GradScaler`, `CosineAnnealingLR`, and `LinearWarmup`.

```python
optimizer = Adam(
                unet.parameters(),
                lr = unet_lr,
                eps = unet_eps,
                betas = (beta1, beta2),
                **kwargs
            )

            if self.use_ema:
                self.ema_unets.append(EMA(unet, **ema_kwargs))

            scaler = GradScaler(enabled = grad_scaler_enabled)

            scheduler = warmup_scheduler = None

            if exists(unet_cosine_decay_max_steps):
                scheduler = CosineAnnealingLR(optimizer, T_max = unet_cosine_decay_max_steps)

            if exists(unet_warmup_steps):
                warmup_scheduler = warmup.LinearWarmup(optimizer, warmup_period = unet_warmup_steps)

                if not exists(scheduler):
                    scheduler = LambdaLR(optimizer, lr_lambda = lambda step: 1.0)

            # set on object

            setattr(self, f'optim{ind}', optimizer) # cannot use pytorch ModuleList for some reason with optimizers
            setattr(self, f'scaler{ind}', scaler)
            setattr(self, f'scheduler{ind}', scheduler)
            setattr(self, f'warmup{ind}', warmup_scheduler)
```

--------------------------------

### Process files with repository metadata

Source: https://context7.com/yilinjz/astchunk/llms.txt

Read a file from disk and process it with repository-level metadata and chunk expansion enabled.

```python
from astchunk import ASTChunkBuilder
import os

chunk_builder = ASTChunkBuilder(
    max_chunk_size=1500,
    language="python",
    metadata_template="default"
)

# Process a single file
input_file = "src/utils/helpers.py"

with open(input_file, 'r', encoding='utf-8') as f:
    code = f.read()

chunks = chunk_builder.chunkify(
    code,
    repo_level_metadata={
        "filepath": input_file
    },
    chunk_expansion=True
)

# Save chunks to separate files
output_dir = "chunks"
os.makedirs(output_dir, exist_ok=True)

for i, chunk in enumerate(chunks, 1):
    output_path = os.path.join(output_dir, f"chunk_{i}.py")
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(chunk['content'])

    print(f"Saved chunk {i}: {chunk['metadata']['chunk_size']} chars")
```

--------------------------------

### Training and Validation Steps

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/fixed_chunking_results.txt

Core execution logic for performing training and validation steps using dataloader iterators.

```python
    def train_step(self, *, unet_number = None, **kwargs):
        if not self.prepared:
            self.prepare()
        self.create_train_iter()

        kwargs = {'unet_number': unet_number, **kwargs}
        loss = self.step_with_dl_iter(self.train_dl_iter, **kwargs)
        self.update(unet_number = unet_number)
        return loss

    @torch.no_grad()
    @eval_decorator
    def valid_step(self, **kwargs):
        if not self.prepared:
            self.prepare()
        self.create_valid_iter()
        context = self.use_ema_unets if kwargs.pop('use_ema_unets', False) else nullcontext
        with context():
            loss = self.step_with_dl_iter(self.valid_dl_iter, **kwargs)
        return loss

    def step_with_dl_iter(self, dl_iter, **kwargs):
        dl_tuple_output = cast_tuple(next(dl_iter))
        model_input = dict(list(zip(self.dl_tuple_output_keywords_names, dl_tuple_output)))
        loss = self.forward(**{**kwargs, **model_input})
        return loss
```

--------------------------------

### Configure Accelerator Scaler

Source: https://github.com/yilinjz/astchunk/blob/main/examples/source_code.txt

Patches the optimizer step method to support per-optimizer gradient scaling within the accelerator.

```python
def set_accelerator_scaler(self, unet_number):
        def patch_optimizer_step(accelerated_optimizer, method):
            def patched_step(*args, **kwargs):
                accelerated_optimizer._accelerate_step_called = True
                return method(*args, **kwargs)
            return patched_step

        unet_number = self.validate_unet_number(unet_number)
        scaler = getattr(self, f'scaler{unet_number - 1}')

        self.accelerator.scaler = scaler
        for optimizer in self.accelerator._optimizers:
            optimizer.scaler = scaler
            optimizer._accelerate_step_called = False
            optimizer._optimizer_original_step_method = optimizer.optimizer.step
            optimizer._optimizer_patched_step_method = patch_optimizer_step(optimizer, optimizer.optimizer.step)
```

--------------------------------

### Processing Multiple Languages with ASTChunkBuilder

Source: https://context7.com/yilinjz/astchunk/llms.txt

Demonstrates creating separate ASTChunkBuilder instances for different programming languages, ensuring correct parsing and chunking for each.

```python
from astchunk import ASTChunkBuilder

# Python builder
python_builder = ASTChunkBuilder(
    max_chunk_size=1500,
    language="python",
    metadata_template="default"
)

```

--------------------------------

### Initialize ASTChunkBuilder for Multiple Languages

Source: https://github.com/yilinjz/astchunk/blob/main/README.md

Shows how to initialize the ASTChunkBuilder for different programming languages like Python, Java, and TypeScript. Each builder can be configured with specific max chunk sizes and metadata templates.

```python
# Python code
python_builder = ASTChunkBuilder(
    max_chunk_size=1500,
    language="python",
    metadata_template="default"
)

# Java code  
java_builder = ASTChunkBuilder(
    max_chunk_size=2000,
    language="java", 
    metadata_template="default"
)

# TypeScript code
ts_builder = ASTChunkBuilder(
    max_chunk_size=1800,
    language="typescript",
    metadata_template="default"
)
```

--------------------------------

### Initialize Gradient Scaler

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/fixed_chunking_results.txt

Initializes GradScaler for mixed-precision training, enabled based on the grad_scaler_enabled flag.

```python
scaler = GradScaler(enabled = grad_scaler_enabled)
```

--------------------------------

### Chunk Expansion with Metadata Headers

Source: https://context7.com/yilinjz/astchunk/llms.txt

Enables chunk expansion to add metadata headers, including file path and hierarchy, to each chunk. Set `chunk_expansion=True` and provide `repo_level_metadata`.

```python
from astchunk import ASTChunkBuilder

chunk_builder = ASTChunkBuilder(
    max_chunk_size=800,
    language="python",
    metadata_template="default"
)

code = """
class ImageProcessor:
    def __init__(self, config):
        self.config = config

    def resize(self, image, width, height):
        # Resize logic here
        return resized_image

    def apply_filter(self, image, filter_name):
        if filter_name == 'blur':
            return self._blur(image)
        return image

    def _blur(self, image):
        # Blur implementation
        return blurred_image
"""

# Enable chunk expansion with file path metadata
chunks = chunk_builder.chunkify(
    code,
    repo_level_metadata={
        "filepath": "src/image/processor.py"
    },
    chunk_expansion=True  # Adds metadata header to each chunk
)

# Output includes header like:
# '''
# src/image/processor.py
# class ImageProcessor:
# '''
# <actual chunk content>

for chunk in chunks:
    print(chunk['content'])
    print("=" * 50)

```

--------------------------------

### Load Imagen Checkpoint

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_with_expansion_results.txt

Loads the Imagen model state, optimizer states, scheduler states, and EMA weights from a specified path. Handles version mismatches and optional loading of only the model.

```python
    def load(self, path, only_model = False, strict = True, noop_if_not_exist = False):
        fs = self.fs

        if noop_if_not_exist and not fs.exists(path):
            self.print(f'trainer checkpoint not found at {str(path)}')
            return

        assert fs.exists(path), f'{path} does not exist'

        self.reset_ema_unets_all_one_device()

        # to avoid extra GPU memory usage in main process when using Accelerate

        with fs.open(path) as f:
            loaded_obj = torch.load(f, map_location='cpu')

        if version.parse(__version__) != version.parse(loaded_obj['version']):
            self.print(f'loading saved imagen at version {loaded_obj["version"]}, but current package version is {__version__}')

        try:
            self.imagen.load_state_dict(loaded_obj['model'], strict = strict)
        except RuntimeError:
            print("Failed loading state dict. Trying partial load")
            self.imagen.load_state_dict(restore_parts(self.imagen.state_dict(),
                                                      loaded_obj['model']))

        if only_model:
            return loaded_obj

        self.steps.copy_(loaded_obj['steps'])

        for ind in range(0, self.num_unets):
            scaler_key = f'scaler{ind}'
            optimizer_key = f'optim{ind}'
            scheduler_key = f'scheduler{ind}'
            warmup_scheduler_key = f'warmup{ind}'

            scaler = getattr(self, scaler_key)
            optimizer = getattr(self, optimizer_key)
            scheduler = getattr(self, scheduler_key)
            warmup_scheduler = getattr(self, warmup_scheduler_key)

            if exists(scheduler) and scheduler_key in loaded_obj:
                scheduler.load_state_dict(loaded_obj[scheduler_key])

            if exists(warmup_scheduler) and warmup_scheduler_key in loaded_obj:
                warmup_scheduler.load_state_dict(loaded_obj[warmup_scheduler_key])

            if exists(optimizer):
                try:
                    optimizer.load_state_dict(loaded_obj[optimizer_key])
                    scaler.load_state_dict(loaded_obj[scaler_key])
                except:
                    self.print('could not load optimizer and scaler, possibly because you have turned on mixed precision training since the last run. resuming with new optimizer and scalers')
```

--------------------------------

### Sample from Model with EMA Unets

Source: https://github.com/yilinjz/astchunk/blob/main/examples/outputs/ast_chunking_results.txt

Generates samples from the model, optionally using Exponential Moving Average (EMA) unets for potentially higher quality results. Handles device placement and progress bar display.

```python
@torch.no_grad()
    @cast_torch_tensor
    @imagen_sample_in_chunks
    def sample(self, *args, **kwargs):
        context = nullcontext if  kwargs.pop('use_non_ema', False) else self.use_ema_unets

        self.print_untrained_unets()

        if not self.is_main:
            kwargs['use_tqdm'] = False

        with context():
            output = self.imagen.sample(*args, device = self.device, **kwargs)

        return output
```
Tokens:
Raw