Update modeling_rqa.py

modeling_rqa.py

@@ -1,8 +1,12 @@
-from typing import Any, Dict, List, Optional
-
 import torch
 import torch.nn as nn
-from transformers import AutoConfig, AutoModel, PreTrainedModel, PretrainedConfig
+from typing import List, Optional
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    PreTrainedModel,
+    PretrainedConfig,
+)
 
 
 class RQAModelConfig(PretrainedConfig):
@@ -11,7 +15,7 @@ class RQAModelConfig(PretrainedConfig):
     def __init__(
         self,
         base_model_name: str = "FacebookAI/xlm-roberta-large",
-        encoder_config: Optional[Dict[str, Any]] = None,
+        encoder_config: Optional[dict] = None,
         error_types: Optional[List[str]] = None,
         schema_version: str = "rqa.v2.2",
         has_issue_projection_dim: int = 256,
@@ -19,7 +23,7 @@ class RQAModelConfig(PretrainedConfig):
         errors_projection_dim: int = 512,
         has_issue_dropout: float = 0.25,
         hidden_dropout: float = 0.25,
-        errors_dropout: float = 0.30,
+        errors_dropout: float = 0.3,
         temperature_has_issue: float = 1.0,
         temperature_is_hidden: float = 1.0,
         temperature_errors: Optional[List[float]] = None,
@@ -27,28 +31,35 @@ class RQAModelConfig(PretrainedConfig):
         threshold_is_hidden: float = 0.5,
         threshold_error: float = 0.5,
         threshold_errors: Optional[List[float]] = None,
-        **kwargs,
+        **kwargs
     ):
         super().__init__(**kwargs)
 
-        self.schema_version = str(schema_version)
         self.base_model_name = base_model_name
         self.encoder_config = encoder_config
-        self.error_types = list(error_types or [])
+        self.error_types = error_types or [
+            "false_causality",
+            "unsupported_claim",
+            "overgeneralization",
+            "missing_premise",
+            "contradiction",
+            "circular_reasoning",
+        ]
         self.num_error_types = len(self.error_types)
+        self.schema_version = schema_version
 
-        self.has_issue_projection_dim = int(has_issue_projection_dim)
-        self.hidden_projection_dim = int(hidden_projection_dim)
-        self.errors_projection_dim = int(errors_projection_dim)
+        self.has_issue_projection_dim = has_issue_projection_dim
+        self.hidden_projection_dim = hidden_projection_dim
+        self.errors_projection_dim = errors_projection_dim
 
-        self.has_issue_dropout = float(has_issue_dropout)
-        self.hidden_dropout = float(hidden_dropout)
-        self.errors_dropout = float(errors_dropout)
+        self.has_issue_dropout = has_issue_dropout
+        self.hidden_dropout = hidden_dropout
+        self.errors_dropout = errors_dropout
 
         self.temperature_has_issue = float(temperature_has_issue)
         self.temperature_is_hidden = float(temperature_is_hidden)
         self.temperature_errors = (
-            list(temperature_errors)
+            temperature_errors
             if temperature_errors is not None
             else [1.0] * self.num_error_types
         )
@@ -57,9 +68,9 @@ class RQAModelConfig(PretrainedConfig):
         self.threshold_is_hidden = float(threshold_is_hidden)
         self.threshold_error = float(threshold_error)
         self.threshold_errors = (
-            list(threshold_errors)
+            threshold_errors
             if threshold_errors is not None
-            else [self.threshold_error] * self.num_error_types
+            else [float(threshold_error)] * self.num_error_types
         )
 
         try:
@@ -69,32 +80,8 @@ class RQAModelConfig(PretrainedConfig):
             pass
 
 
-def build_encoder_config_from_saved_dict(
-    encoder_config: Optional[Dict[str, Any]],
-    base_model_name: str,
-):
-    if encoder_config is None:
-        return AutoConfig.from_pretrained(base_model_name)
-
-    cfg_dict = dict(encoder_config)
-    model_type = cfg_dict.pop("model_type", None)
-    cfg_dict.pop("_name_or_path", None)
-
-    if model_type is not None:
-        try:
-            return AutoConfig.for_model(model_type, **cfg_dict)
-        except Exception:
-            pass
-
-    return AutoConfig.from_pretrained(base_model_name)
-
-
 class MeanPooling(nn.Module):
-    def forward(
-        self,
-        last_hidden_state: torch.Tensor,
-        attention_mask: torch.Tensor,
-    ) -> torch.Tensor:
+    def forward(self, last_hidden_state, attention_mask):
         mask = attention_mask.unsqueeze(-1).float()
         summed = torch.sum(last_hidden_state * mask, dim=1)
         denom = torch.clamp(mask.sum(dim=1), min=1e-9)
@@ -106,24 +93,17 @@ class RQAModelHF(PreTrainedModel):
     _supports_grouped_mm = False
 
     def __init__(self, config: RQAModelConfig):
+        super().__init__(config)
+
         try:
             config._experts_implementation = "eager"
             config._experts_implementation_internal = "eager"
         except Exception:
             pass
-        super().__init__(config)
-
-        if config.encoder_config is None:
-            base_cfg = AutoConfig.from_pretrained(config.base_model_name)
-            config.encoder_config = base_cfg.to_dict()
-
-        enc_cfg = build_encoder_config_from_saved_dict(
-            encoder_config=config.encoder_config,
-            base_model_name=config.base_model_name,
-        )
-        self.encoder = AutoModel.from_config(enc_cfg)
 
+        self.encoder = AutoModel.from_pretrained(config.base_model_name)
         hidden_size = self.encoder.config.hidden_size
+
         self.pooler = MeanPooling()
 
         self.has_issue_projection = nn.Sequential(
@@ -132,12 +112,14 @@ class RQAModelHF(PreTrainedModel):
             nn.GELU(),
             nn.Dropout(config.has_issue_dropout),
         )
+
         self.hidden_projection = nn.Sequential(
             nn.Linear(hidden_size, config.hidden_projection_dim),
             nn.LayerNorm(config.hidden_projection_dim),
             nn.GELU(),
             nn.Dropout(config.hidden_dropout),
         )
+
         self.errors_projection = nn.Sequential(
             nn.Linear(hidden_size, config.errors_projection_dim),
             nn.LayerNorm(config.errors_projection_dim),
@@ -155,11 +137,10 @@ class RQAModelHF(PreTrainedModel):
         self.log_var_has_issue = nn.Parameter(torch.zeros(1))
         self.log_var_is_hidden = nn.Parameter(torch.zeros(1))
         self.log_var_errors = nn.Parameter(torch.zeros(1))
-        with torch.no_grad():
-            self.log_var_has_issue.clamp_(-5, 5)
-            self.log_var_is_hidden.clamp_(-5, 5)
-            self.log_var_errors.clamp_(-5, 5)
 
+        self._init_custom_weights()
+
+    def _init_custom_weights(self):
         for module in [
             self.has_issue_projection[0],
             self.hidden_projection[0],
@@ -168,47 +149,40 @@ class RQAModelHF(PreTrainedModel):
             self.is_hidden_head,
             self.errors_head,
         ]:
-            setattr(module, "_rqa_custom_init", True)
-
-        self.post_init()
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
 
-    def _init_weights(self, module):
-        if isinstance(module, nn.Linear) and getattr(module, "_rqa_custom_init", False):
-            nn.init.xavier_uniform_(module.weight)
-            if module.bias is not None:
-                nn.init.zeros_(module.bias)
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
-        **kwargs,
-    ) -> Dict[str, torch.Tensor]:
+    def forward(self, input_ids=None, attention_mask=None, **kwargs):
         outputs = self.encoder(
             input_ids=input_ids,
             attention_mask=attention_mask,
             return_dict=True,
-            **kwargs,
         )
+
         pooled = self.pooler(outputs.last_hidden_state, attention_mask)
 
-        issue_features = self.has_issue_projection(pooled)
-        hidden_features = self.hidden_projection(pooled)
-        error_features = self.errors_projection(pooled)
+        has_issue_logits = self.has_issue_head(
+            self.has_issue_projection(pooled)
+        ).squeeze(-1)
+
+        is_hidden_logits = self.is_hidden_head(
+            self.hidden_projection(pooled)
+        ).squeeze(-1)
+
+        errors_logits = self.errors_head(
+            self.errors_projection(pooled)
+        )
 
         return {
-            "has_issue_logits": self.has_issue_head(issue_features).squeeze(-1),
-            "is_hidden_logits": self.is_hidden_head(hidden_features).squeeze(-1),
-            "errors_logits": self.errors_head(error_features),
+            "has_issue_logits": has_issue_logits,
+            "is_hidden_logits": is_hidden_logits,
+            "errors_logits": errors_logits,
         }
 
 
-try:
-    AutoConfig.register("rqa_v2_2", RQAModelConfig)
-except ValueError:
-    pass
+AutoConfig.register("rqa_v2_2", RQAModelConfig)
+AutoModel.register(RQAModelConfig, RQAModelHF)
 
-try:
-    AutoModel.register(RQAModelConfig, RQAModelHF)
-except ValueError:
-    pass
+print("✅ RQA-R2 зарегистрирован в Transformers")