| 28 | OPTILLM_MODEL_NAME = "codelion/optillm-modernbert-large" |
| 29 | |
| 30 | class OptILMClassifier(nn.Module): |
| 31 | def __init__(self, base_model, num_labels): |
| 32 | super().__init__() |
| 33 | self.base_model = base_model |
| 34 | self.effort_encoder = nn.Sequential( |
| 35 | nn.Linear(1, 64), |
| 36 | nn.ReLU(), |
| 37 | nn.Linear(64, 64), |
| 38 | nn.ReLU() |
| 39 | ) |
| 40 | self.classifier = nn.Linear(base_model.config.hidden_size + 64, num_labels) |
| 41 | |
| 42 | def forward(self, input_ids, attention_mask, effort): |
| 43 | outputs = self.base_model(input_ids=input_ids, attention_mask=attention_mask) |
| 44 | pooled_output = outputs.last_hidden_state[:, 0] # Shape: (batch_size, hidden_size) |
| 45 | effort_encoded = self.effort_encoder(effort.unsqueeze(1)) # Shape: (batch_size, 64) |
| 46 | combined_input = torch.cat((pooled_output, effort_encoded), dim=1) |
| 47 | logits = self.classifier(combined_input) |
| 48 | return logits |
| 49 | |
| 50 | def load_optillm_model(): |
| 51 | device = torch.device("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu") |
no outgoing calls
no test coverage detected