(args)
| 233 | return results |
| 234 | |
| 235 | def main(args): |
| 236 | |
| 237 | if args.push_to_hub: |
| 238 | base_model = AutoModel.from_pretrained(args.model_name) |
| 239 | tokenizer = AutoTokenizer.from_pretrained(args.model_name) |
| 240 | # best_model = OptILMClassifier(base_model, num_labels=len(APPROACHES)) |
| 241 | # best_model.to(device) |
| 242 | # load_model(best_model, "best_model.safetensors") |
| 243 | # we just push the base model and then upload the safetensors file manually as OptILMClassifier class doesn't have a push_to_hub method. |
| 244 | base_model.push_to_hub(args.hub_model_id) |
| 245 | tokenizer.push_to_hub(args.hub_model_id) |
| 246 | return |
| 247 | |
| 248 | tokenizer = AutoTokenizer.from_pretrained(args.model_name) |
| 249 | dataset = load_and_preprocess_data(tokenizer) |
| 250 | |
| 251 | kf = KFold(n_splits=args.k_folds, shuffle=True, random_state=42) |
| 252 | |
| 253 | best_val_accuracy = 0 |
| 254 | best_fold = 0 |
| 255 | |
| 256 | for fold, (train_indices, val_indices) in enumerate(kf.split(dataset), 1): |
| 257 | print(f"\nTraining Fold {fold}") |
| 258 | |
| 259 | train_sampler = SubsetRandomSampler(train_indices) |
| 260 | val_sampler = SubsetRandomSampler(val_indices) |
| 261 | |
| 262 | train_dataloader = DataLoader(dataset, batch_size=args.batch_size, sampler=train_sampler) |
| 263 | val_dataloader = DataLoader(dataset, batch_size=args.batch_size, sampler=val_sampler) |
| 264 | |
| 265 | base_model = AutoModel.from_pretrained(args.model_name) |
| 266 | model = OptILMClassifier(base_model, num_labels=len(APPROACHES)).to(device) |
| 267 | |
| 268 | optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=0.01) |
| 269 | scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.1, patience=2, verbose=True) |
| 270 | |
| 271 | train(model, train_dataloader, val_dataloader, optimizer, scheduler, args.num_epochs, args.patience, args.clip_value) |
| 272 | |
| 273 | # Evaluate the model on the validation set |
| 274 | fold_val_accuracy = validate(model, val_dataloader) |
| 275 | print(f"Fold {fold} Validation Accuracy: {fold_val_accuracy:.4f}") |
| 276 | |
| 277 | # Save the model for this fold |
| 278 | save_model(model, f"model_fold_{fold}.safetensors") |
| 279 | |
| 280 | # Update best model if this fold performed better |
| 281 | if fold_val_accuracy > best_val_accuracy: |
| 282 | best_val_accuracy = fold_val_accuracy |
| 283 | best_fold = fold |
| 284 | save_model(model, "best_model.safetensors") |
| 285 | |
| 286 | print(f"\nBest performing model was from fold {best_fold} with validation accuracy {best_val_accuracy:.4f}") |
| 287 | |
| 288 | # Load the best model for inference |
| 289 | base_model = AutoModel.from_pretrained(args.model_name) |
| 290 | best_model = OptILMClassifier(base_model, num_labels=len(APPROACHES)) |
| 291 | best_model.to(device) |
| 292 | load_model(best_model, "best_model.safetensors") |
no test coverage detected