|
|
import torch |
|
|
import torch.nn as nn |
|
|
from transformers import AutoModel |
|
|
|
|
|
class MultitaskCodeSimilarityModel(nn.Module): |
|
|
def __init__(self, config, tokenizer): |
|
|
super().__init__() |
|
|
self.config = config |
|
|
self.tokenizer = tokenizer |
|
|
self.encoder = AutoModel.from_config(config) |
|
|
self.classifier = nn.Linear(config.hidden_size, config.num_labels) |
|
|
|
|
|
|
|
|
self.decoder_embedding = nn.Linear(config.hidden_size, config.hidden_size) |
|
|
self.decoder = nn.GRU( |
|
|
input_size=config.hidden_size, |
|
|
hidden_size=config.hidden_size, |
|
|
batch_first=True |
|
|
) |
|
|
self.explanation_head = nn.Linear(config.hidden_size, len(tokenizer)) |
|
|
|
|
|
def forward(self, input_ids, attention_mask, explanation_ids=None, explanation_mask=None): |
|
|
outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask) |
|
|
pooled = outputs.last_hidden_state[:, 0] |
|
|
logits = self.classifier(pooled) |
|
|
|
|
|
explanation_logits = None |
|
|
if explanation_ids is not None: |
|
|
decoder_input = self.decoder_embedding(pooled).unsqueeze(1).expand(-1, explanation_ids.size(1), -1) |
|
|
decoder_outputs, _ = self.decoder(decoder_input) |
|
|
explanation_logits = self.explanation_head(decoder_outputs) |
|
|
|
|
|
return logits, explanation_logits |
