Metadata-Version: 2.4
Name: celmoe-vp
Version: 3.0.0
Summary: Composable Expert Layer MoE primitives with hierarchical routing, fusion, and multi-loss orchestration.
Author: F000NK, Voluntas Progressus
License-Expression: MIT
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.14
Classifier: Operating System :: OS Independent
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0.0
Dynamic: license-file

# celmoe-vp

`celmoe-vp` — доменно-агностичная библиотека для **иерархических Mixture-of-Experts**: явная маршрутизация по уровням, объединение представлений уровней (fusion) и **составная целевая функция** (global / per-level / bridge + лоссы экспертов).

PyPI:

```bash
pip install celmoe-vp
```

Импорт в Python:

```python
import celmoe
```

Пакет не привязан к морфологии; это переиспользуемый слой оркестрации для приложений вроде `morphoformer`.

**Документация по устройству пакета:** [ARCHITECTURE.md](ARCHITECTURE.md) (MoE + иерархия + multi-task, схема модулей, mermaid-поток данных).

---

## Что такое CELMoE здесь

**CELMoE** трактуется как три механизма в одном API:

| Компонент | Смысл |
|-----------|--------|
| **MoE** | На каждом уровне — набор экспертов; для каждого элемента батча выбирается **имя эксперта** (список длины `B` на уровень). |
| **Иерархия** | Уровни в `CELMoEConfig.levels` выполняются **по порядку**; скрытое состояние передаётся «вниз» по стеку уровней. |
| **Multi-task** | Итоговый лосс собирается из нескольких **областей**: глобальные функции, функции на конкретном уровне, **межуровневые (bridge)** и тензоры из `ExpertResult.losses`. |

Вы реализуете содержимое экспертов и при необходимости подменяете модуль fusion.

---

## Структура пакета (v2)

```
celmoe/
  config/       # ExpertConfig, HierarchyLevelConfig, CELMoEConfig
  core/         # типы, ExpertModule, LearnedFusion, HierarchicalCELMoE
  loss/         # LossBundle, CELMoELossAPI
  utils/        # routing_constant, validate_routing_length, merge_routing
```

Стабильный вход — корневой namespace:

```python
from celmoe import CELMoEConfig, HierarchicalCELMoE, CELMoELossAPI, routing_constant
```

Точечные импорты (по желанию):

```python
from celmoe.core.types import LevelRouting, ExpertContext
from celmoe.loss import LossBundle
```

---

## Конфигурация

Публичный слой — dataclass-ы:

- **`ExpertConfig`** — один эксперт: `loss_weight`, `stop_gradient_from_parent`, произвольный `metadata`.
- **`HierarchyLevelConfig`** — уровень (`universal`, `family`, `language`, …): словарь экспертов, опциональный `fallback_expert`, флаг `enabled`.
- **`CELMoEConfig`** — `hidden_size`, список уровней, параметры fusion (`fusion_dropout`, `use_learned_fusion` зарезервировано для совместимости; фактическая подмена fusion — через аргумент конструктора модели).

---

## Контракт эксперта

Эксперты наследуют **`ExpertModule`**:

```python
from celmoe import ExpertBatch, ExpertModule, ExpertResult

class MyExpert(ExpertModule):
    def forward(self, batch: ExpertBatch) -> ExpertResult:
        return ExpertResult(hidden=batch.hidden)
```

**`ExpertBatch`:** `hidden`, `context`, `mask`, `state` (строковые метки вроде уровня/эксперта).

**`ExpertResult`:** `hidden`, словарь `losses` (попадут в общий `CELMoEOutput` с префиксами), произвольный `aux`.

---

## Иерархическое выполнение

**`HierarchicalCELMoE`** прогоняет включённые уровни, для каждого уровня применяет выбранных по маршруту экспертов, собирает лоссы уровня и передаёт скрытое состояние дальше. Затем все `LevelOutput.hidden` объединяются через **fusion** (по умолчанию **`LearnedFusion`**; можно передать свой `nn.Module`).

Формат маршрута — **`LevelRouting`**: `dict[str, list[str]]`, каждый список длины батча:

```python
{
    "universal": ["core", "core", "core"],
    "family": ["slavic", "romance", "slavic"],
    "language": ["rus", "spa", "bul"],
}
```

Если на уровне ровно один эксперт или задан **`fallback_expert`**, недостающий ключ маршрута может быть заполнен автоматически (broadcast).

Вспомогательные функции:

```python
from celmoe import routing_constant, validate_routing_length, merge_routing

routes = routing_constant("core", batch_size=8)
validate_routing_length("universal", routes, 8)
combined = merge_routing(existing_routing, {"new_level": ["a"] * 8})
```

---

## Multi-loss / multi-task API

**`CELMoELossAPI`** регистрирует:

- **`add_global`** — лосс от всего `CELMoEOutput`;
- **`add_level`** — лосс от `LevelOutput` выбранного уровня;
- **`add_bridge`** — лосс от пары родитель/дочерний уровень;
- **`compute`** — собирает **`LossBundle`** (взвешенная сумма, плюс `as_dict()` для логов).

Уже на этапе `compute` в бандл попадают лоссы из `output.losses` (эксперты) с scope `"expert"`. У каждого **`LossTerm`** есть строка **`scope`** для фильтрации или отображения в логгерах.

Это позволяет комбинировать, например: регуляризацию родителя, специализацию ребёнка, согласование между уровнями — без жёсткой привязки к одной задаче внутри библиотеки.

---

## Пример

```python
from typing import Mapping

import torch
import torch.nn as nn

from celmoe import (
    CELMoEConfig,
    CELMoELossAPI,
    ExpertBatch,
    ExpertConfig,
    ExpertModule,
    ExpertResult,
    HierarchicalCELMoE,
    HierarchyLevelConfig,
)


class LinearExpert(ExpertModule):
    def __init__(self, hidden_size: int) -> None:
        super().__init__()
        self.proj = nn.Linear(hidden_size, hidden_size)

    def forward(self, batch: ExpertBatch) -> ExpertResult:
        hidden = self.proj(batch.hidden)
        return ExpertResult(hidden=hidden)


def build_expert(level_name: str, expert_name: str, metadata: Mapping[str, object]) -> ExpertModule:
    del level_name, expert_name, metadata
    return LinearExpert(hidden_size=256)


config = CELMoEConfig(
    hidden_size=256,
    levels=[
        HierarchyLevelConfig(
            name="parent",
            experts={"core": ExpertConfig(name="core")},
            fallback_expert="core",
        ),
        HierarchyLevelConfig(
            name="child",
            experts={
                "a": ExpertConfig(name="a", stop_gradient_from_parent=True),
                "b": ExpertConfig(name="b", stop_gradient_from_parent=True),
            },
        ),
    ],
)

model = HierarchicalCELMoE(config, expert_factory=build_expert)

hidden = torch.randn(4, 32, 256)
output = model(
    hidden,
    routing={
        "parent": ["core"] * 4,
        "child": ["a", "b", "a", "b"],
    },
)

loss_api = CELMoELossAPI()
loss_api.add_global("l2", lambda out, targets, context: out.fused_hidden.pow(2).mean())
bundle = loss_api.compute(output)
print(bundle.total)
```

---

## Основные типы данных

- **`LevelOutput`** — результат одного уровня: скрытое состояние, маршрут, лоссы, `aux`.
- **`CELMoEOutput`** — все уровни, `fused_hidden`, агрегированные лоссы экспертов.
- **`LossTerm` / `LossBundle`** — пошаговая разложимость лосса и скаляры для метрик.

---

## Изоляция градиента

**`stop_gradient_from_parent=True`** у **`ExpertConfig`** отключает градиент от входа уровня к выбранному эксперту (через `detach` входа). Полезно для поэтапного обучения и снижения интерференции между уровнями.

---

## Что пакет намеренно не включает

- эмбеддинги и токенизацию;
- внутренности Transformer-блоков;
- загрузку данных, оптимизаторы, чекпоинты.

Это остаётся на стороне приложения; `celmoe-vp` — слой оркестрации MoE + иерархии + составного лосса.

---

## Типизация

Пакет строго типизирован и поставляет **`py.typed`** для использования как независимого зависимого пакета.

---

## Версия 2.0

В **2.0** изменена **внутренняя** раскладка модулей (`config` / `core` / `loss` / `utils`); публичный импорт `from celmoe import ...` сохранён и расширен (типы, `LearnedFusion`, утилиты маршрутизации). Подробности — в [ARCHITECTURE.md](ARCHITECTURE.md).
