|
|
|
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
from typing import List, Optional |
|
|
from transformers import ( |
|
|
AutoModel, |
|
|
PreTrainedModel, |
|
|
PretrainedConfig, |
|
|
AutoConfig, |
|
|
AutoModel, |
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class RQAModelConfig(PretrainedConfig): |
|
|
model_type = "rqa" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
base_model_name: str = "FacebookAI/xlm-roberta-large", |
|
|
num_error_types: int = 6, |
|
|
has_issue_projection_dim: int = 256, |
|
|
errors_projection_dim: int = 512, |
|
|
has_issue_dropout: float = 0.25, |
|
|
errors_dropout: float = 0.3, |
|
|
temperature_has_issue: float = 1.0, |
|
|
temperature_errors: Optional[List[float]] = None, |
|
|
**kwargs |
|
|
): |
|
|
super().__init__(**kwargs) |
|
|
|
|
|
self.base_model_name = base_model_name |
|
|
self.num_error_types = num_error_types |
|
|
self.has_issue_projection_dim = has_issue_projection_dim |
|
|
self.errors_projection_dim = errors_projection_dim |
|
|
self.has_issue_dropout = has_issue_dropout |
|
|
self.errors_dropout = errors_dropout |
|
|
|
|
|
self.temperature_has_issue = temperature_has_issue |
|
|
self.temperature_errors = ( |
|
|
temperature_errors |
|
|
if temperature_errors is not None |
|
|
else [1.0] * num_error_types |
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class MeanPooling(nn.Module): |
|
|
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) |
|
|
return summed / denom |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class RQAModelHF(PreTrainedModel): |
|
|
config_class = RQAModelConfig |
|
|
|
|
|
def __init__(self, config: RQAModelConfig): |
|
|
super().__init__(config) |
|
|
|
|
|
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( |
|
|
nn.Linear(hidden_size, config.has_issue_projection_dim), |
|
|
nn.LayerNorm(config.has_issue_projection_dim), |
|
|
nn.GELU(), |
|
|
nn.Dropout(config.has_issue_dropout), |
|
|
) |
|
|
|
|
|
self.errors_projection = nn.Sequential( |
|
|
nn.Linear(hidden_size, config.errors_projection_dim), |
|
|
nn.LayerNorm(config.errors_projection_dim), |
|
|
nn.GELU(), |
|
|
nn.Dropout(config.errors_dropout), |
|
|
) |
|
|
|
|
|
self.has_issue_head = nn.Linear(config.has_issue_projection_dim, 1) |
|
|
self.errors_head = nn.Linear( |
|
|
config.errors_projection_dim, config.num_error_types |
|
|
) |
|
|
|
|
|
self._init_custom_weights() |
|
|
|
|
|
def _init_custom_weights(self): |
|
|
for module in [ |
|
|
self.has_issue_projection[0], |
|
|
self.errors_projection[0], |
|
|
self.has_issue_head, |
|
|
self.errors_head, |
|
|
]: |
|
|
if isinstance(module, nn.Linear): |
|
|
nn.init.xavier_uniform_(module.weight) |
|
|
nn.init.zeros_(module.bias) |
|
|
|
|
|
def forward(self, input_ids=None, attention_mask=None, **kwargs): |
|
|
outputs = self.encoder( |
|
|
input_ids=input_ids, |
|
|
attention_mask=attention_mask, |
|
|
return_dict=True, |
|
|
) |
|
|
|
|
|
pooled = self.pooler(outputs.last_hidden_state, attention_mask) |
|
|
|
|
|
has_issue_logits = self.has_issue_head( |
|
|
self.has_issue_projection(pooled) |
|
|
).squeeze(-1) |
|
|
|
|
|
errors_logits = self.errors_head( |
|
|
self.errors_projection(pooled) |
|
|
) |
|
|
|
|
|
return { |
|
|
"has_issue_logits": has_issue_logits, |
|
|
"errors_logits": errors_logits, |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
AutoConfig.register("rqa", RQAModelConfig) |
|
|
AutoModel.register(RQAModelConfig, RQAModelHF) |
|
|
|
|
|
print("✅ RQA зарегистрирован в Transformers") |
