cida-plugin 1.0.0__py3-none-any.whl

cida_plugin/hf.py

@@ -0,0 +1,194 @@
+"""
+hf.py — Интеграция CIDA-Plugin с Hugging Face Models.
+
+Позволяет обернуть любую модель классификации последовательностей (Sequence Classification) 
+в CIDA-Plugin в одну строчку кода, заменяя стандартную линейную голову.
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Union, Dict, Any
+
+try:
+    from transformers.modeling_outputs import SequenceClassifierOutput
+    _HAS_TRANSFORMERS = True
+except ImportError:
+    # Заглушка на случай отсутствия transformers (например, при тестировании голого PyTorch)
+    class SequenceClassifierOutput:
+        def __init__(self, loss=None, logits=None, hidden_states=None, attentions=None):
+            self.loss = loss
+            self.logits = logits
+            self.hidden_states = hidden_states
+            self.attentions = attentions
+    _HAS_TRANSFORMERS = False
+
+from .core import CIDAPlugin
+from .config import CIDAPluginConfig
+from .losses import CIDALoss
+
+
+class HFModelWithCIDA(nn.Module):
+    """
+    Обертка над предобученной Hugging Face моделью классификации текстов.
+    Перехватывает скрытые состояния базового энкодера и пропускает их через CIDA-Plugin.
+    """
+    def __init__(
+        self,
+        hf_model: nn.Module,
+        cida_config: Optional[CIDAPluginConfig] = None,
+        **cida_kwargs: Any
+    ):
+        super().__init__()
+        
+        # Сохраняем исходную конфигурацию
+        self.config = getattr(hf_model, "config", None)
+        
+        # 1. Извлекаем базовый энкодер (transformer backbone)
+        if hasattr(hf_model, "base_model"):
+            self.encoder = hf_model.base_model
+        else:
+            # Fallback для кастомных моделей
+            self.encoder = hf_model
+            
+        # 2. Определяем размерности
+        if self.config is not None:
+            d_model = getattr(self.config, "hidden_size", None) or getattr(self.config, "d_model", None)
+            num_labels = getattr(self.config, "num_labels", 2)
+        else:
+            d_model = cida_kwargs.get("d_input")
+            num_labels = cida_kwargs.get("num_classes", 2)
+            
+        if d_model is None:
+            raise ValueError(
+                "Не удалось автоматически определить hidden_size из конфигурации модели HF. "
+                "Пожалуйста, передайте параметр d_input явно в cida_kwargs."
+            )
+            
+        # 3. Инициализируем CIDA-Plugin
+        if cida_config is None:
+            cida_config = CIDAPluginConfig(
+                d_input=d_model,
+                num_classes=num_labels,
+                **cida_kwargs
+            )
+        else:
+            # Обновляем поля, если они переданы пустые
+            if cida_config.d_input is None:
+                cida_config.d_input = d_model
+            if cida_config.num_classes is None:
+                cida_config.num_classes = num_labels
+                
+        self.plugin = CIDAPlugin(cida_config)
+        self.num_labels = num_labels
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs: Any
+    ) -> Union[SequenceClassifierOutput, tuple]:
+        
+        use_return_dict = return_dict if return_dict is not None else getattr(self.config, "use_return_dict", True)
+        
+        # 1. Прогоняем через базовый энкодер
+        encoder_outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=use_return_dict,
+            **kwargs
+        )
+        
+        # 2. Извлекаем sequence outputs
+        if isinstance(encoder_outputs, tuple):
+            last_hidden_state = encoder_outputs[0]
+        else:
+            last_hidden_state = encoder_outputs.last_hidden_state
+            
+        # Извлекаем CLS токен (первый токен последовательности) как pooled representation
+        pooled_output = last_hidden_state[:, 0, :]
+        
+        # 3. Вызываем CIDA-Plugin
+        plugin_outputs = self.plugin(
+            pooled=pooled_output,
+            seq_output=last_hidden_state,
+            mask=attention_mask
+        )
+        
+        logits = plugin_outputs["p_final"]
+        
+        # 4. Вычисляем лосс при наличии меток
+        loss = None
+        if labels is not None:
+            loss_fn = CIDALoss(
+                multi_label=self.plugin.config.multi_label,
+                lambda_cal=self.plugin.config.lambda_cal,
+                lambda_ac=self.plugin.config.lambda_ac,
+                min_disagreement=self.plugin.config.min_disagreement
+            ).to(logits.device)
+            
+            orth_loss = plugin_outputs.get("L_orth", None)
+            
+            loss, _ = loss_fn(
+                p_final=logits,
+                y=labels,
+                b_all=plugin_outputs["b_all"],
+                orth_loss=orth_loss
+            )
+            
+        # 5. Возвращаем результат
+        if not use_return_dict:
+            output = (logits,)
+            if isinstance(encoder_outputs, tuple):
+                output = output + encoder_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+            
+        # Собираем SequenceClassifierOutput для HF-совместимости
+        res = SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=getattr(encoder_outputs, "hidden_states", None),
+            attentions=getattr(encoder_outputs, "attentions", None)
+        )
+        
+        # Добавляем CIDA метаданные как атрибуты для диагностики
+        res.b_all = plugin_outputs["b_all"]
+        res.uncertainty = plugin_outputs["uncertainty"]
+        res.disagreement = plugin_outputs["disagreement"]
+        res.rounds_used = plugin_outputs["rounds_used"]
+        res.plugin_outputs = plugin_outputs
+        
+        return res
+
+
+def wrap_hf_model(
+    hf_model: nn.Module,
+    cida_config: Optional[CIDAPluginConfig] = None,
+    **kwargs: Any
+) -> HFModelWithCIDA:
+    """
+    Быстрая обертка для Hugging Face sequence classification моделей.
+    
+    Пример использования:
+        from transformers import AutoModelForSequenceClassification
+        from cida_plugin import wrap_hf_model
+
+        model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)
+        model = wrap_hf_model(model, max_rounds=3)
+        
+        # Теперь модель готова к обучению с CIDA-Plugin в Hugging Face Trainer!
+    """
+    return HFModelWithCIDA(hf_model, cida_config=cida_config, **kwargs)

cida_plugin/liquid_dynamics.py

@@ -0,0 +1,291 @@
+"""
+Liquid Neural ODE Dynamics для CIDA-Plugin.
+
+Заменяет дискретный цикл deliberation loop на непрерывное
+дифференциальное уравнение:
+    ds/dt = -s/τ(x) + F(s, r, e)
+
+Где:
+    s - скрытое состояние агентов
+    r - контр-аргументы от других агентов
+    e - извлечённые доказательства
+    τ(x) - зависящая от входа константа времени (Liquid Time Constant)
+    F - функция перехода (основанная на AgentUpdater)
+
+Использование `torchdiffeq` позволяет использовать адаптивные решатели
+(например, dopri5), которые автоматически регулируют количество
+вычислений (NFE) в зависимости от сложности входа.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from torchdiffeq import odeint_adjoint as odeint_adj
+    from torchdiffeq import odeint
+    _HAS_ADJOINT = True
+except ImportError:
+    from torchdiffeq import odeint
+    odeint_adj = odeint
+    _HAS_ADJOINT = False
+
+from .agent import AgentState
+from .deliberation import (
+    AgentEvidenceExtractor,
+    CounterargumentCommunication,
+    AgentUpdater,
+)
+
+
+class LiquidTimeConstant(nn.Module):
+    """
+    Вычисляет τ(x) — константу времени, зависящую от начального состояния.
+    """
+
+    def __init__(self, d_hidden: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(d_hidden, d_hidden // 2),
+            nn.SiLU(),
+            nn.Linear(d_hidden // 2, 1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, s_init: torch.Tensor) -> torch.Tensor:
+        """
+        s_init: (B, M, d)
+        Returns: (B, M, 1) — значения τ ∈ (0.1, 1.1)
+        """
+        return self.net(s_init) + 0.1
+
+
+class LiquidDeliberationDynamics(nn.Module):
+    """
+    Обёртка над компонентами deliberation (extractor, communication, updater),
+    которая вычисляет производную ds/dt для ODE решателя.
+
+    Ключевое отличие от предыдущей версии: beliefs (b) и uncertainty (u)
+    вычисляются ДИНАМИЧЕСКИ из текущего скрытого состояния s через
+    updater.g(s), а не подставляются как нулевые заглушки.
+    """
+
+    def __init__(
+        self,
+        extractor: AgentEvidenceExtractor,
+        message_formulator,
+        communication: CounterargumentCommunication,
+        updater: AgentUpdater,
+        d_hidden: int,
+        trajectory_save_every: int = 10,
+        T: float = 4.0,
+    ):
+        super().__init__()
+        self.extractor = extractor
+        self.message_formulator = message_formulator
+        self.communication = communication
+        self.updater = updater
+        self.tau_net = LiquidTimeConstant(d_hidden)
+        self.trajectory_save_every = trajectory_save_every
+        self.T = T
+        # Временные точки для сохранения (корректно для adaptive solvers)
+        self._save_timepoints = {0.25, 0.50, 0.75, 1.0}
+
+        # Кэшируемый контекст (чтобы не прокидывать через ODE)
+        self.memory = None
+        self.mask = None
+        self.tau = None
+        self.nfe = 0  # Number of Function Evaluations
+
+        # Для сохранения промежуточных состояний (trajectory)
+        self.b_trajectory = []
+        self.u_trajectory = []
+        self.p_ptr_trajectory = []
+        self.m_trajectory = []
+
+    def set_context(
+        self,
+        memory: torch.Tensor,
+        mask: torch.Tensor,
+        s_init: torch.Tensor,
+    ):
+        """
+        Устанавливает контекст перед запуском ODE.
+        memory: (B, L, d) — память документов
+        mask: (B, L) — маска внимания
+        s_init: (B, M, d) — начальное состояние агентов
+        """
+        self.memory = memory
+        self.mask = mask
+        self.nfe = 0
+        # FIX v4: del [:] вместо clear() для явного освобождения памяти
+        del self.b_trajectory[:]
+        del self.u_trajectory[:]
+        del self.p_ptr_trajectory[:]
+        del self.m_trajectory[:]
+
+    def clear_context(self):
+        self.memory = None
+        self.mask = None
+
+    def _compute_beliefs(self, s: torch.Tensor):
+        """
+        Динамически вычисляет текущие beliefs и uncertainty из состояния s.
+
+        Вместо передачи dummy_b=0, dummy_u=1 (что лишало коммуникацию
+        информации о текущих мнениях агентов), мы проецируем s через
+        ту же голову, что использует AgentUpdater.
+
+        Returns: (current_b, current_u)
+        """
+        g_out = self.updater.g(s)
+
+        if self.updater.multi_label:
+            current_b = torch.sigmoid(g_out)
+            current_u = (4.0 * current_b * (1.0 - current_b)).mean(
+                dim=-1, keepdim=True
+            )
+        else:
+            alpha = F.softplus(g_out) + 1.0
+            alpha_sum = alpha.sum(dim=-1, keepdim=True)
+            current_b = alpha / alpha_sum
+            current_u = float(g_out.size(-1)) / alpha_sum
+
+        return current_b, current_u
+
+    def forward(self, t: torch.Tensor, s: torch.Tensor) -> torch.Tensor:
+        """
+        Вычисляет ds/dt.
+        t: (1,) — текущее время ODE
+        s: (B, M, d) — текущее состояние агентов
+        """
+        self.nfe += 1
+
+        B, M = s.shape[0], s.shape[1]
+        num_classes = self.updater.g[-1].out_features
+
+        # Динамическое вычисление текущих beliefs и uncertainty
+        current_b, current_u = self._compute_beliefs(s)
+
+        # Evidence extraction
+        p_ptr_t, e_t = self.extractor(s, self.memory, self.mask)
+
+        # Communication с актуальными beliefs (не с нулями!)
+        m_t = self.message_formulator(s, e_t, current_b, current_u)
+        r_t = self.communication(s, m_t, current_b, e_t)
+
+        # AgentUpdater step
+        dummy_state = AgentState(
+            s=s,
+            b=current_b,
+            u=current_u,
+            p=torch.zeros(B, M, 1, device=s.device),
+            e=e_t,
+            alpha=torch.ones(B, M, num_classes, device=s.device),
+        )
+        next_state = self.updater(dummy_state, r_t, e_t)
+
+        # Сохраняем промежуточные beliefs для лоссов и консенсуса
+        # FIX v4: 1) detach() — без него ODE граф держится в памяти (сотни МБ)
+        #         2) temporal checkpoints вместо NFE-based (корректно для adaptive solvers)
+        if torch.is_grad_enabled() or not self.training:
+            t_normalized = t.item() / self.T if self.T > 0 else 0
+            should_save = any(abs(t_normalized - sp) < 0.02 for sp in self._save_timepoints)
+            if should_save or self.nfe % self.trajectory_save_every == 0:
+                self.b_trajectory.append(next_state.b.detach())
+                self.u_trajectory.append(next_state.u.detach())
+                self.p_ptr_trajectory.append(p_ptr_t.detach())
+                self.m_trajectory.append(s.detach())
+
+        # TRUE Liquid ODE v4: τ(s,t) вычисляется на КАЖДОМ шаге из текущего s
+        # (раньше τ фиксировался из s_init в set_context — это не настоящий Liquid NN)
+        tau = self.tau_net(s)  # (B, M, 1) — state-dependent time constant
+        f_s = next_state.s - s
+        ds_dt = -s / tau + f_s
+
+        return ds_dt
+
+
+class LiquidDeliberationSolver(nn.Module):
+    """
+    Запускает ODE решатель для deliberation.
+    
+    При use_adjoint=True использует adjoint method для экономии VRAM:
+    градиенты считаются без хранения промежуточных состояний.
+    """
+
+    def __init__(
+        self,
+        extractor: AgentEvidenceExtractor,
+        message_formulator,
+        communication: CounterargumentCommunication,
+        updater: AgentUpdater,
+        d_hidden: int,
+        solver: str = "euler",
+        atol: float = 1e-3,
+        rtol: float = 1e-3,
+        T: float = 4.0,
+        trajectory_save_every: int = 10,
+        use_adjoint: bool = False,
+    ):
+        super().__init__()
+        self.dynamics = LiquidDeliberationDynamics(
+            extractor=extractor,
+            message_formulator=message_formulator,
+            communication=communication,
+            updater=updater,
+            d_hidden=d_hidden,
+            trajectory_save_every=trajectory_save_every,
+            T=T,
+        )
+        self.solver = solver
+        self.atol = atol
+        self.rtol = rtol
+        self.T = T
+        self.use_adjoint = use_adjoint and _HAS_ADJOINT
+        self._integrate = odeint_adj if self.use_adjoint else odeint
+
+    def forward(
+        self,
+        s_init: torch.Tensor,
+        memory: torch.Tensor,
+        mask: torch.Tensor,
+    ) -> tuple:
+        """
+        s_init: (B, M, d)
+        memory: (B, L, d)
+        mask: (B, L)
+
+        Returns:
+            s_final (B, M, d)
+            trajectory (dict)
+        """
+        device = s_init.device
+        self.dynamics.set_context(memory, mask, s_init)
+
+        t = torch.tensor([0.0, self.T], device=device)
+
+        states = self._integrate(
+            self.dynamics,
+            s_init,
+            t,
+            method=self.solver,
+            atol=self.atol,
+            rtol=self.rtol,
+        )
+
+        s_final = states[-1]
+        nfe = self.dynamics.nfe
+
+        trajectory = {
+            "b_all": self.dynamics.b_trajectory,
+            "u_all": self.dynamics.u_trajectory,
+            "p_ptr_all": self.dynamics.p_ptr_trajectory,
+            "m_all": self.dynamics.m_trajectory,
+            "nfe": nfe,
+        }
+
+        self.dynamics.clear_context()
+
+        return s_final, trajectory
+

cida_plugin/losses.py

@@ -0,0 +1,174 @@
+"""
+losses.py — Упрощённая система потерь CIDA (3 компонента вместо 11)
+
+Было (OmegaLossSystem):
+    task + calibration + debate + progress + budget + dominance
+    + role + role_spec + orthogonality + ponder_cost + ttt_reg
+    = 11 компонентов, 11 гиперпараметров
+
+Стало (CIDALoss):
+    task + calibration + anti_collapse
+    = 3 компонента, 2 гиперпараметра
+
+Почему можно удалить остальные 8:
+
+    debate_loss      → удалён: несогласие теперь структурное (role priors)
+    progress_loss    → удалён: избыточен с task_loss при правильных priors
+    budget_loss      → удалён: ACT не нужен если число раундов мало (≤3)
+    dominance_loss   → удалён: симптом проблемы с PoE, которую мы исправили
+    role_loss        → удалён: нет RoleEmbeddings → нет и orth-лосса для них
+    role_spec_loss   → удалён: роли теперь structural, не statistical
+    orth_loss        → заменён одной строкой в compute_role_orthogonality_loss()
+    ttt_reg          → удалён: TTT убран из базовой архитектуры
+
+anti_collapse_loss:
+    Единственный "новый" компонент, но он ПРОЩЕ и ПРАВИЛЬНЕЕ debate_loss.
+    Вместо "попади в целевой уровень несогласия на каждом раунде" →
+    "просто не позволяй агентам стать идентичными".
+    Это МИНИМАЛЬНОЕ требование, а не жёсткое целевое расписание.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class CIDALoss(nn.Module):
+    """
+    Упрощённая система потерь для CIDA-Plugin.
+
+    Параметры
+    ---------
+    lambda_cal : float
+        Вес calibration loss (Brier score). Рекомендуется 0.3–0.5.
+    lambda_ac : float
+        Вес anti-collapse loss. Рекомендуется 0.1–0.3.
+        Меньше → агенты могут сближаться (риск коллапса).
+        Больше → агенты разведены сильнее (риск потери точности).
+    min_disagreement : float
+        Минимальный порог несогласия. Рекомендуется 0.05–0.15.
+        При role priors реальное несогласие ≥ 0.21, поэтому loss
+        почти всегда = 0 (действует как safety net).
+    multi_label : bool
+        Если True — используется BCE вместо CE.
+    """
+
+    def __init__(
+        self,
+        lambda_cal: float = 0.4,
+        lambda_ac: float = 0.2,
+        min_disagreement: float = 0.08,
+        multi_label: bool = False,
+    ):
+        super().__init__()
+        self.lambda_cal = lambda_cal
+        self.lambda_ac = lambda_ac
+        self.min_disagreement = min_disagreement
+        self.multi_label = multi_label
+
+    # ── Компонент 1: Task Loss ────────────────────────────────────────────────
+
+    def task_loss(self, p_final: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        """
+        Основная задача: классификация.
+        """
+        eps = 1e-7
+        if self.multi_label:
+            return F.binary_cross_entropy(p_final.clamp(eps, 1.0 - eps), y)
+        # FIX v4: log-softmax стабильнее чем clamp().log() — согласовано с consensus.py
+        log_probs = F.log_softmax(p_final.log().clamp(-10, 10), dim=-1)
+        return F.nll_loss(log_probs, y)
+
+    # ── Компонент 2: Calibration Loss ─────────────────────────────────────────
+
+    def calibration_loss(self, p_final: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        """
+        Brier score: MSE между вероятностями и метками.
+        """
+        if self.multi_label:
+            probs = p_final
+            return ((probs - y) ** 2).mean()
+
+        # Single-label: one-hot targets
+        K = p_final.size(-1)
+        y_oh = F.one_hot(y, num_classes=K).float()
+        return ((p_final - y_oh) ** 2).mean()
+
+    # ── Компонент 3: Anti-Collapse Loss ──────────────────────────────────────
+
+    def anti_collapse_loss(self, b_final: torch.Tensor) -> torch.Tensor:
+        """
+        Предотвращает коллапс агентов к идентичным убеждениям.
+
+        НЕ является аналогом debate_loss. Ключевые отличия:
+
+        debate_loss:
+            - Целевое расписание d_t на каждый раунд
+            - Штрафует если несогласие != цели
+            - Создаёт противоречие с consensus
+
+        anti_collapse_loss:
+            - Только минимальная граница min_disagreement
+            - Штрафует только при несогласии НИЖЕ порога
+            - Нулевой градиент при нормальном обучении (safety net)
+
+        b_final: (B, M, K) — финальные убеждения агентов
+        Returns: scalar (обычно ≈ 0 если role priors работают)
+
+        Математика:
+            d_ij = ||b_i - b_j||_1 — несогласие между агентами i и j
+            L_ac = mean(relu(δ - d_ij)) где δ = min_disagreement
+            Градиент ненулевой только когда d_ij < δ
+        """
+        b_i = b_final.unsqueeze(2)   # (B, M, 1, K)
+        b_j = b_final.unsqueeze(1)   # (B, 1, M, K)
+        pairwise_dist = (b_i - b_j).abs().sum(dim=-1)  # (B, M, M)
+
+        # Только нижняя треугольная часть (без диагонали)
+        M = b_final.size(1)
+        mask = torch.tril(torch.ones(M, M, device=b_final.device), diagonal=-1).bool()
+        off_diag = pairwise_dist[:, mask]  # (B, n_pairs)
+
+        return F.relu(self.min_disagreement - off_diag).mean()
+
+    # ── Основной forward ──────────────────────────────────────────────────────
+
+    def forward(
+        self,
+        p_final: torch.Tensor,
+        y: torch.Tensor,
+        b_all: list,
+        orth_loss: torch.Tensor = None,
+    ) -> tuple[torch.Tensor, dict]:
+        """
+        p_final : (B, K)         — финальное предсказание
+        y       : (B,) или (B,K) — метки
+        b_all   : list[(B,M,K)]  — история убеждений по раундам
+        orth_loss: scalar или None — orth loss из representations (опционально)
+
+        Returns: (total_loss, components_dict)
+        """
+        l_t = self.task_loss(p_final, y)
+        l_c = self.calibration_loss(p_final, y)
+
+        # Anti-collapse на финальных убеждениях
+        b_final = b_all[-1] if b_all else None
+        l_ac = self.anti_collapse_loss(b_final) if b_final is not None else torch.tensor(0.0)
+
+        total = l_t + self.lambda_cal * l_c + self.lambda_ac * l_ac
+
+        # Orth loss опционально (из representations, не из весов)
+        l_orth = torch.tensor(0.0)
+        if orth_loss is not None:
+            l_orth = orth_loss
+            total = total + 0.05 * l_orth  # малый вес, это regularizer
+
+        components = {
+            "task":         l_t.item(),
+            "calibration":  l_c.item(),
+            "anti_collapse": l_ac.item(),
+            "orth":         l_orth.item() if isinstance(l_orth, torch.Tensor) else l_orth,
+            "total":        total.item(),
+        }
+
+        return total, components