Add the documentation for package include API, Examples and guides.

Author: codewithdark-git

CHANGELOG.md

@@ -0,0 +1,38 @@
+# Changelog
+
+All notable changes to DiffusionLM will be documented in this file.
+
+## [0.1.6] - 2025-04-05
+
+### Added
+- Initial release of DiffusionLM
+- Basic transformer-based diffusion model
+- Training pipeline with error handling
+- Dataset preparation utilities
+- Documentation and examples
+- Multi-GPU training support
+- Mixed precision training
+- Custom generation strategies
+
+### Fixed
+- Memory optimization for large models
+- Import path resolution
+- Package structure and dependencies
+
+## [released]
+
+### Added
+- Streaming token generation
+- Advanced logging configuration
+- Performance optimization guides
+- Comprehensive documentation
+
+### Changed
+- Improved error handling
+- Updated package structure
+- Enhanced docstrings
+
+### Fixed
+- Package import issues
+- Memory leaks in training
+- Documentation build process

docs/Makefile

@@ -0,0 +1,15 @@
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = .
+BUILDDIR      = _build
+
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile clean
+
+clean:
+	rm -rf $(BUILDDIR)/*
+
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

docs/api/model.md

@@ -0,0 +1,91 @@
+# Model API Reference
+
+## DiffusionConfig
+
+Configuration class for the DiffusionLM model.
+
+### Parameters
+
+- `vocab_size` (int, default=50257): Size of the vocabulary
+- `hidden_size` (int, default=768): Dimensionality of the hidden layers
+- `num_hidden_layers` (int, default=12): Number of transformer layers
+- `num_attention_heads` (int, default=12): Number of attention heads per layer
+- `intermediate_size` (int, default=3072): Dimensionality of feed-forward layers
+- `hidden_dropout_prob` (float, default=0.1): Dropout probability for hidden layers
+- `attention_probs_dropout_prob` (float, default=0.1): Dropout probability for attention
+- `max_position_embeddings` (int, default=1024): Maximum sequence length
+- `num_timesteps` (int, default=100): Number of diffusion timesteps
+- `time_embed_dim` (int, default=128): Dimensionality of time embeddings
+
+## DiffusionLLM
+
+Main model class implementing the transformer-based diffusion language model.
+
+### Methods
+
+#### forward()
+
+```python
+def forward(
+    input_ids=None,
+    attention_mask=None,
+    timesteps=None,
+    labels=None,
+    return_dict=True
+)
+```
+
+Perform a forward pass through the model.
+
+#### generate()
+
+```python
+def generate(
+    prompt=None,
+    max_length=100,
+    num_inference_steps=50,
+    temperature=1.0,
+    strategy='random',
+    top_p=0.9,
+    top_k=50,
+    num_beams=5,
+    return_scores=False,
+    use_streaming=False,
+    callback_fn=None
+)
+```
+
+Generate text using the reverse diffusion process.
+
+## MultiHeadAttention
+
+Multi-head self-attention mechanism implementation.
+
+### Methods
+
+#### forward()
+
+```python
+def forward(
+    hidden_states,
+    attention_mask=None,
+    head_mask=None,
+    output_attentions=False
+)
+```
+
+Perform multi-head attention computation.
+
+## TimeEmbedding
+
+Embedding layer for diffusion timesteps.
+
+### Methods
+
+#### forward()
+
+```python
+def forward(timesteps)
+```
+
+Generate time embeddings for the given timesteps.

docs/api/trainer.md

@@ -0,0 +1,63 @@
+# Trainer API Reference
+
+## Trainer Function
+
+Main training function for DiffusionLM models.
+
+```python
+def trainer(
+    model: DiffusionLLM,
+    train_dataset,
+    val_dataset = None,
+    batch_size: int = 8,
+    num_epochs: int = 5,
+    learning_rate: float = 5e-5,
+    warmup_steps: int = 1000,
+    max_grad_norm: float = 1.0,
+    num_timesteps: int = 100,
+    save_path: Optional[str] = None,
+    device: torch.device = None
+) -> DiffusionLLM
+```
+
+### Parameters
+
+- `model`: The DiffusionLLM model to train
+- `train_dataset`: Training dataset
+- `val_dataset`: Validation dataset (optional)
+- `batch_size`: Batch size for training
+- `num_epochs`: Number of training epochs
+- `learning_rate`: Learning rate
+- `warmup_steps`: Number of warmup steps
+- `max_grad_norm`: Maximum gradient norm
+- `num_timesteps`: Number of diffusion timesteps
+- `save_path`: Path to save checkpoints
+- `device`: Device to train on
+
+### Returns
+
+- Trained DiffusionLLM model
+
+## Evaluate Function
+
+```python
+def evaluate(
+    model: DiffusionLLM,
+    dataloader: DataLoader,
+    device: torch.device,
+    num_timesteps: int = 100,
+    num_eval_steps: int = None
+) -> float
+```
+
+### Parameters
+
+- `model`: Model to evaluate
+- `dataloader`: DataLoader for evaluation data
+- `device`: Device to evaluate on
+- `num_timesteps`: Number of diffusion timesteps
+- `num_eval_steps`: Number of evaluation steps
+
+### Returns
+
+- Average loss on the evaluation set

docs/api/utils.md

@@ -0,0 +1,67 @@
+# Utils API Reference
+
+## Dataset Preparation
+
+### prepare_dataset()
+
+```python
+def prepare_dataset(
+    dataset_name: str = "wikitext/wikitext-2-raw-v1",
+    tokenizer_name: str = "gpt2",
+    max_length: int = 1024,
+    cache_dir: Optional[str] = None,
+    num_proc: int = 4
+) -> Tuple[PYTORCH_Dataset, Optional[PYTORCH_Dataset], AutoTokenizer]
+```
+
+Prepares datasets for training DiffusionLM models.
+
+### PYTORCH_Dataset
+
+```python
+class PYTORCH_Dataset(Dataset):
+    def __init__(
+        self,
+        dataset: Any,
+        mask_token_id: int,
+        pad_token_id: int,
+    )
+```
+
+Custom dataset class for DiffusionLM training.
+
+## Error Handling
+
+### DiffusionLMError
+
+```python
+class DiffusionLMError(Exception):
+    """Base exception class for DiffusionLM package"""
+    pass
+```
+
+### handle_errors()
+
+```python
+def handle_errors(
+    error_class: Type[Exception] = DiffusionLMError, 
+    reraise: bool = True, 
+    logger: logging.Logger = logger
+) -> Callable
+```
+
+Decorator for handling errors in functions.
+
+## Logging
+
+### setup_logging()
+
+```python
+def setup_logging(
+    log_file: str = None, 
+    level: int = logging.INFO,
+    format: str = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+) -> None
+```
+
+Sets up logging configuration for the package.

docs/build_docs.md

@@ -0,0 +1,61 @@
+# Building the Documentation
+
+## Requirements
+
+First, install the required packages:
+
+```bash
+pip install -r docs/requirements.txt
+```
+
+## Building HTML Documentation
+
+### On Linux/Mac:
+```bash
+cd docs
+make html
+```
+
+### On Windows:
+```bash
+cd docs
+./make.bat html
+```
+
+The built documentation will be available in `docs/_build/html/`.
+
+## Development Server
+
+For live preview while writing documentation:
+
+```bash
+pip install sphinx-autobuild
+sphinx-autobuild docs docs/_build/html
+```
+
+Then visit `http://localhost:8000` in your browser.
+
+## Build Options
+
+- `make html` - Build HTML documentation
+- `make clean` - Clean build directory
+- `make help` - Show all available build options
+
+## Troubleshooting
+
+If you encounter any issues:
+
+1. Ensure all requirements are installed:
+   ```bash
+   pip install -r docs/requirements.txt
+   ```
+
+2. Clean the build directory:
+   ```bash
+   make clean
+   ```
+
+3. Check for syntax errors in rst/md files
+   ```bash
+   sphinx-build -nW -b html docs/ docs/_build/html
+   ```

docs/conf.py

@@ -0,0 +1,27 @@
+import os
+import sys
+sys.path.insert(0, os.path.abspath('..'))
+
+project = 'DiffusionLM'
+copyright = '2024, Dark Coder'
+author = 'Dark Coder'
+release = '0.1.0'
+
+extensions = [
+    'sphinx.ext.autodoc',
+    'sphinx.ext.napoleon',
+    'sphinx.ext.viewcode',
+    'sphinx.ext.githubpages',
+    'myst_parser',
+]
+
+templates_path = ['_templates']
+exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
+
+html_theme = 'sphinx_rtd_theme'
+html_static_path = ['_static']
+
+autodoc_member_order = 'bysource'
+autodoc_typehints = 'description'
+napoleon_google_docstring = True
+napoleon_numpy_docstring = True