Feature/emotional optimizer

.gitignore

@@ -205,3 +205,6 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+# Data from examples
+data/

README.md

@@ -1,96 +1,152 @@
-# EmotiGrad
+# 🌈 EmotiGrad — Emotional Support for Your Optimizers
+
+<p align="center">
+  <img src="https://img.shields.io/badge/status-pre--release-blueviolet" />
+  <img src="https://img.shields.io/badge/tests-passing-brightgreen" />
+  <img src="https://img.shields.io/badge/code%20style-black-000000" />
+  <img src="https://img.shields.io/badge/linting-ruff-8A2BE2" />
+  <img src="https://img.shields.io/badge/license-MIT-green" />
+</p>
 
 EmotiGrad is a tiny Python library that wraps your PyTorch optimizers and gives you emotionally-charged feedback during training, from wholesome encouragement to unhinged sass.
 
 It aims to be:
-- **Drop-in friendly** – keep your usual `torch.optim` code
-- **Fun but useful** – emotional logs + basic training insights
-- **Extensible** – easily add new "personalities" and behaviors
-
-> ### Because sometimes you need more than just `.step()`, you need support.
 
+* **Drop-in friendly**: swap it into any `torch.optim` workflow
+* **Fun but useful**: emotional logs + basic training insights
+* **Extensible**: easily add new "personalities" and behaviors
 
 ## Status
 
-> ⚠️ EmotiGrad is under active early development (pre-release).  
-> The API may change before `0.1.0`. Feedback and ideas are very welcome!
+> ⚠️ EmotiGrad is under active early development (pre-release).
+> Expect the API to evolve before version `0.1.0`.
+> Feedback and ideas are *very* welcome!
 
 ---
 
 ## Installation
 
-For now, install from source:
+Install from source:
 
 ```bash
 git clone git@github.com:smiley-maker/emotigrad.git
 cd emotigrad
 pip install -e .
 ```
 
-(PyPI support will come in a later release.)
+PyPI packages will come in a later release.
+
+---
 
 ## Quick Start
 
-Basic Usage: 
+Here’s the smallest possible example:
 
 ```python
 import torch
 from emotigrad import EmotionalOptimizer
 
 model = torch.nn.Linear(10, 1)
-base_optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+base_opt = torch.optim.Adam(model.parameters(), lr=1e-3)
 
 # Wrap your optimizer with a personality
-optimizer = EmotionalOptimizer(
-    base_optimizer,
-    personality="wholesome",  # (planned) "sassy", "chaotic", etc.
+opt = EmotionalOptimizer(
+    base_opt,
+    personality="wholesome",   # also: "sassy", "quiet", custom callables, etc.
+    message_every=20,           # feedback every 20 steps (averaged)
 )
 
-for step in range(10):
+for step in range(50):
     x = torch.randn(32, 10)
     y = torch.randn(32, 1)
 
     preds = model(x)
     loss = (preds - y).pow(2).mean()
 
-    optimizer.optimizer.zero_grad()
+    opt.zero_grad()
     loss.backward()
 
-    # In a future version, this call will emit emotional messages
-    optimizer.step()
+    # Provide loss to trigger feedback
+    opt.step(loss=loss.item())
+```
+
+### How `message_every` works
+
+Instead of reacting to every single (noisy) loss value, EmotiGrad:
+
+1. Collects the last **N** loss values
+2. Computes the **average loss for that block**
+3. Compares it to the **previous block’s average**
+4. Feeds the result into your chosen personality
+
+This produces smoother, more meaningful emotional feedback.
+
+Set `message_every=1` for per-step chatter.
+
+
+## Personalities
+
+EmotiGrad ships with several built-in personalities, such as:
+
+* **wholesome** – kind, encouraging, proud of your progress
+* **sassy** – mildly offended by your gradients
+* **quiet** – reports occasionally, like a stoic mentor
+
+You can also write your own:
+
+```python
+def hype(loss, prev, step):
+    if prev and loss < prev:
+        return f"🚀 Step {step}: HUGE gains! {prev:.4f} → {loss:.4f}"
+    return None
+
+opt = EmotionalOptimizer(base_opt, personality=hype)
+```
+
+Or register them globally:
+
+```python
+from emotigrad.personalities import register_personality
+
+register_personality("hype", hype)
+opt = EmotionalOptimizer(base_opt, personality="hype")
 ```
 
-In upcoming versions, `optimizer.step(loss=loss.item())` will trigger personality-specific messages based on how training is going.
+## Examples
 
-## Roadmap (high level)
+You can find examples in the `examples/` directory:
+
+* `basic_usage.py`
+* `mnist_training.py`
+* `custom_personality.py`
+
+## Roadmap
 
 Planned features:
 
-- Emotional personas: 
-    - wholesome – positive, encouraging
-    - sassy – mildly offended by your gradients
-    - chaotic – unhelpful but entertaining
-- Basic training trend detection (loss going up/down, plateauing)
-- Configurable verbosity and logging destinations
-- Easy hooks for custom personalities
-- Longer-term ideas:
-    - Integration with PyTorch Lightning / HuggingFace Trainer
-    - Optional LLM-based "training advisor" for suggestions
+* More built-in emotional personas:
+  * `wholesome`, `sassy`, `quiet`, `chaotic`, `roaster`, `nervous`, etc.
+* Trend-aware training feedback
+* Configurable output formatting (e.g. text colors and formatting)
+* Easy hooks for custom personalities
+* Optional LLM-based “training advisor” mode
+* Integrations with:
+  * PyTorch Lightning
+  * HuggingFace Trainer
+* Rich visual outputs (ASCII art, emoji graphs, etc.)
 
-## Contributions
+## Contributing
 
-Contributions are very welcome, even at this early stage!
+Contributions are warmly welcome — even small improvements help!
 
-Some helpful ways to contribute:
+Ways to contribute:
 
-- Try EmotiGrad in a toy project and open issues for:
-    - bugs
-    - confusing APIs
-    - personality ideas
-    - Add or improve a personality preset
-- Add tests or docs
+* Report bugs or confusing APIs
+* Suggest new personalities or features
+* Improve tests or documentation
+* Add examples
 
-To setup for development: 
+To set up development:
 
 ```bash
 git clone git@github.com:smiley-maker/emotigrad.git
@@ -99,23 +155,33 @@ pip install -e ".[dev]"
 pytest
 ```
 
-We also suggest using a virtual environment or conda to manage dependencies. Development dependencies and more detailed instructions will come as the project evolves. 
+We also recommend using a virtual environment or conda during local development.
 
-## Current Project Structure
+
+## Project Structure
 
 ```
 emotigrad/
-  emotigrad/
-    __init__.py
-    base.py           # will hold EmotionalOptimizer
+  src/
+    emotigrad/
+      __init__.py
+      base.py           # EmotionalOptimizer
+      personalities.py  # built-in personas + registry
+      types.py          # Personality Protocol
   tests/
-    test_smoke.py     # tiny test so CI has something
+  examples/
   README.md
   LICENSE
   pyproject.toml
-  .gitignore
 ```
 
+
 ## License
 
-EmotiGrad is open source and under the MIT License. Learn more in the LICENSE file. 
+EmotiGrad is released under the MIT License.
+See the LICENSE file for details.
+
+
+## Thanks for checking out EmotiGrad
+
+If you build something with it, please share it or open an issue, I’d love to see what you make!

examples/basic_usage.py

@@ -0,0 +1,50 @@
+"""
+Basic usage example for EmotiGrad.
+
+This script shows how to wrap a PyTorch optimizer with an EmotionalOptimizer
+and get emotionally-enhanced training feedback. It uses a tiny synthetic dataset
+so it runs instantly and without external dependencies.
+
+Run with:
+    python examples/basic_usage.py
+"""
+
+import torch
+
+from emotigrad import EmotionalOptimizer
+
+
+def main():
+    # Simple linear model for demonstration
+    model = torch.nn.Linear(10, 1)
+
+    # Standard optimizer
+    base_opt = torch.optim.Adam(model.parameters(), lr=1e-3)
+
+    # Wrap with EmotiGrad!
+    opt = EmotionalOptimizer(
+        base_opt,
+        personality="wholesome",  # try "sassy" or write your own!
+        message_every=5,  # emotional feedback every 5 steps
+    )
+
+    # Synthetic training loop
+    for step in range(20):
+        x = torch.randn(32, 10)
+        y = torch.randn(32, 1)
+
+        preds = model(x)
+        loss = (preds - y).pow(2).mean()
+
+        opt.zero_grad()
+        loss.backward()
+
+        # Passing loss triggers EmotiGrad's emotional feedback
+        opt.step(loss=loss.item())
+
+        if step % 5 == 0:
+            print(f"[step {step}] loss = {loss.item():.4f}")
+
+
+if __name__ == "__main__":
+    main()

examples/custom_personality.py

@@ -0,0 +1,58 @@
+"""
+Example: Creating a custom 'roast' personality for EmotiGrad.
+
+Shows how to write a Personality callable and pass it to EmotionalOptimizer.
+"""
+
+import torch
+
+from emotigrad import EmotionalOptimizer
+
+# --- Custom Personality -------------------------------------------------------
+
+
+def roast(loss, prev_loss, step):
+    """A sarcastic personality that roasts the model's progress."""
+    if prev_loss is None:
+        return f"🔥 Step {step}: New model? Cute. Let's watch it struggle. (avg loss {loss:.4f})"
+
+    if loss < prev_loss:
+        return f"😏 Step {step}: Look at you improving! Honestly shocked. ({prev_loss:.4f} → {loss:.4f})"
+
+    if loss > prev_loss:
+        return (
+            f"🙃 Step {step}: Nice, you made it worse. "
+            f"({prev_loss:.4f} → {loss:.4f}). Truly groundbreaking."
+        )
+
+    return f"🤨 Step {step}: No change. Riveting."
+
+
+# --- Training Loop ------------------------------------------------------------
+
+
+def main():
+    model = torch.nn.Linear(5, 1)
+    base_opt = torch.optim.SGD(model.parameters(), lr=0.1)
+
+    # Use the custom roast personality
+    opt = EmotionalOptimizer(
+        base_opt,
+        personality=roast,  # <-- pass the callable directly
+        message_every=3,  # roast based on averaged loss every 3 steps
+    )
+
+    for step in range(12):
+        x = torch.randn(32, 5)
+        y = torch.randn(32, 1)
+
+        preds = model(x)
+        loss = (preds - y).pow(2).mean()
+
+        opt.zero_grad()
+        loss.backward()
+        opt.step(loss=loss.item())
+
+
+if __name__ == "__main__":
+    main()