(model, train_dataloader, val_dataloader, optimizer, scheduler, num_epochs, patience, clip_value)
| 128 | return logits |
| 129 | |
| 130 | def train(model, train_dataloader, val_dataloader, optimizer, scheduler, num_epochs, patience, clip_value): |
| 131 | best_val_accuracy = 0.0 |
| 132 | epochs_without_improvement = 0 |
| 133 | |
| 134 | for epoch in range(num_epochs): |
| 135 | model.train() |
| 136 | total_loss = 0 |
| 137 | total_accuracy = 0 |
| 138 | |
| 139 | for batch in tqdm(train_dataloader, desc=f"Epoch {epoch+1}/{num_epochs}"): |
| 140 | input_ids = batch['input_ids'].to(device) |
| 141 | attention_mask = batch['attention_mask'].to(device) |
| 142 | approaches = batch['approaches'].to(device) |
| 143 | ranks = batch['ranks'].to(device) |
| 144 | tokens = batch['tokens'].to(device) |
| 145 | |
| 146 | # Normalize tokens to [0, 1] range as a proxy for effort |
| 147 | effort = (tokens - tokens.min()) / (tokens.max() - tokens.min()) |
| 148 | |
| 149 | # Use the minimum rank (best approach) for each prompt |
| 150 | best_approach_indices = ranks.argmin(dim=1) |
| 151 | |
| 152 | logits = model(input_ids, attention_mask, effort[:, 0]) # Use effort for the best approach |
| 153 | |
| 154 | # Calculate standard cross-entropy loss |
| 155 | ce_loss = F.cross_entropy(logits, best_approach_indices) |
| 156 | |
| 157 | # Calculate effort-sensitive loss |
| 158 | effort_loss = F.mse_loss(logits.softmax(dim=1).gather(1, best_approach_indices.unsqueeze(1)).squeeze(), effort[:, 0]) |
| 159 | |
| 160 | # Combine losses |
| 161 | loss = ce_loss + 0.1 * effort_loss # Adjust the weight of effort_loss as needed |
| 162 | |
| 163 | loss.backward() |
| 164 | torch.nn.utils.clip_grad_norm_(model.parameters(), clip_value) |
| 165 | optimizer.step() |
| 166 | optimizer.zero_grad() |
| 167 | |
| 168 | total_loss += loss.item() |
| 169 | predictions = torch.argmax(logits, dim=-1) |
| 170 | total_accuracy += calculate_accuracy(predictions, best_approach_indices) |
| 171 | |
| 172 | avg_train_loss = total_loss / len(train_dataloader) |
| 173 | avg_train_accuracy = total_accuracy / len(train_dataloader) |
| 174 | |
| 175 | # Validation |
| 176 | avg_val_accuracy = validate(model, val_dataloader) |
| 177 | |
| 178 | print(f"Epoch {epoch+1}/{num_epochs}, Train Loss: {avg_train_loss:.4f}, Train Accuracy: {avg_train_accuracy:.4f}, Val Accuracy: {avg_val_accuracy:.4f}") |
| 179 | |
| 180 | # Learning rate scheduling |
| 181 | if isinstance(scheduler, ReduceLROnPlateau): |
| 182 | scheduler.step(avg_val_accuracy) |
| 183 | else: |
| 184 | scheduler.step() |
| 185 | |
| 186 | if avg_val_accuracy > best_val_accuracy: |
| 187 | best_val_accuracy = avg_val_accuracy |
no test coverage detected