skatzR/RQA-R2

1import torch2import torch.nn as nn3from typing import List, Optional4from transformers import (5    AutoConfig,6    AutoModel,7    PreTrainedModel,8    PretrainedConfig,9)10 11 12class RQAModelConfig(PretrainedConfig):13    model_type = "rqa_v2_2"14 15    def __init__(16        self,17        base_model_name: str = "FacebookAI/xlm-roberta-large",18        encoder_config: Optional[dict] = None,19        error_types: Optional[List[str]] = None,20        schema_version: str = "rqa.v2.2",21        has_issue_projection_dim: int = 256,22        hidden_projection_dim: int = 256,23        errors_projection_dim: int = 512,24        has_issue_dropout: float = 0.25,25        hidden_dropout: float = 0.25,26        errors_dropout: float = 0.3,27        temperature_has_issue: float = 1.0,28        temperature_is_hidden: float = 1.0,29        temperature_errors: Optional[List[float]] = None,30        threshold_has_issue: float = 0.5,31        threshold_is_hidden: float = 0.5,32        threshold_error: float = 0.5,33        threshold_errors: Optional[List[float]] = None,34        **kwargs35    ):36        super().__init__(**kwargs)37 38        self.base_model_name = base_model_name39        self.encoder_config = encoder_config40        self.error_types = error_types or [41            "false_causality",42            "unsupported_claim",43            "overgeneralization",44            "missing_premise",45            "contradiction",46            "circular_reasoning",47        ]48        self.num_error_types = len(self.error_types)49        self.schema_version = schema_version50 51        self.has_issue_projection_dim = has_issue_projection_dim52        self.hidden_projection_dim = hidden_projection_dim53        self.errors_projection_dim = errors_projection_dim54 55        self.has_issue_dropout = has_issue_dropout56        self.hidden_dropout = hidden_dropout57        self.errors_dropout = errors_dropout58 59        self.temperature_has_issue = float(temperature_has_issue)60        self.temperature_is_hidden = float(temperature_is_hidden)61        self.temperature_errors = (62            temperature_errors63            if temperature_errors is not None64            else [1.0] * self.num_error_types65        )66 67        self.threshold_has_issue = float(threshold_has_issue)68        self.threshold_is_hidden = float(threshold_is_hidden)69        self.threshold_error = float(threshold_error)70        self.threshold_errors = (71            threshold_errors72            if threshold_errors is not None73            else [float(threshold_error)] * self.num_error_types74        )75 76        try:77            self._experts_implementation = "eager"78            self._experts_implementation_internal = "eager"79        except Exception:80            pass81 82 83class MeanPooling(nn.Module):84    def forward(self, last_hidden_state, attention_mask):85        mask = attention_mask.unsqueeze(-1).float()86        summed = torch.sum(last_hidden_state * mask, dim=1)87        denom = torch.clamp(mask.sum(dim=1), min=1e-9)88        return summed / denom89 90 91class RQAModelHF(PreTrainedModel):92    config_class = RQAModelConfig93    _supports_grouped_mm = False94 95    def __init__(self, config: RQAModelConfig):96        super().__init__(config)97 98        try:99            config._experts_implementation = "eager"100            config._experts_implementation_internal = "eager"101        except Exception:102            pass103 104        self.encoder = AutoModel.from_pretrained(config.base_model_name)105        hidden_size = self.encoder.config.hidden_size106 107        self.pooler = MeanPooling()108 109        self.has_issue_projection = nn.Sequential(110            nn.Linear(hidden_size, config.has_issue_projection_dim),111            nn.LayerNorm(config.has_issue_projection_dim),112            nn.GELU(),113            nn.Dropout(config.has_issue_dropout),114        )115 116        self.hidden_projection = nn.Sequential(117            nn.Linear(hidden_size, config.hidden_projection_dim),118            nn.LayerNorm(config.hidden_projection_dim),119            nn.GELU(),120            nn.Dropout(config.hidden_dropout),121        )122 123        self.errors_projection = nn.Sequential(124            nn.Linear(hidden_size, config.errors_projection_dim),125            nn.LayerNorm(config.errors_projection_dim),126            nn.GELU(),127            nn.Dropout(config.errors_dropout),128        )129 130        self.has_issue_head = nn.Linear(config.has_issue_projection_dim, 1)131        self.is_hidden_head = nn.Linear(config.hidden_projection_dim, 1)132        self.errors_head = nn.Linear(133            config.errors_projection_dim,134            config.num_error_types,135        )136 137        self.log_var_has_issue = nn.Parameter(torch.zeros(1))138        self.log_var_is_hidden = nn.Parameter(torch.zeros(1))139        self.log_var_errors = nn.Parameter(torch.zeros(1))140 141        self._init_custom_weights()142 143    def _init_custom_weights(self):144        for module in [145            self.has_issue_projection[0],146            self.hidden_projection[0],147            self.errors_projection[0],148            self.has_issue_head,149            self.is_hidden_head,150            self.errors_head,151        ]:152            if isinstance(module, nn.Linear):153                nn.init.xavier_uniform_(module.weight)154                if module.bias is not None:155                    nn.init.zeros_(module.bias)156 157    def forward(self, input_ids=None, attention_mask=None, **kwargs):158        outputs = self.encoder(159            input_ids=input_ids,160            attention_mask=attention_mask,161            return_dict=True,162        )163 164        pooled = self.pooler(outputs.last_hidden_state, attention_mask)165 166        has_issue_logits = self.has_issue_head(167            self.has_issue_projection(pooled)168        ).squeeze(-1)169 170        is_hidden_logits = self.is_hidden_head(171            self.hidden_projection(pooled)172        ).squeeze(-1)173 174        errors_logits = self.errors_head(175            self.errors_projection(pooled)176        )177 178        return {179            "has_issue_logits": has_issue_logits,180            "is_hidden_logits": is_hidden_logits,181            "errors_logits": errors_logits,182        }183 184 185AutoConfig.register("rqa_v2_2", RQAModelConfig)186AutoModel.register(RQAModelConfig, RQAModelHF)187 188print("✅ RQA-R2 зарегистрирован в Transformers")189