| 208 | |
| 209 | |
| 210 | def build_model(cfg: DictConfig): |
| 211 | # Note: cfg.num_labels should match the number of classes in your dataset! |
| 212 | if cfg.name == "hf_bert": |
| 213 | return hf_bert_module.create_hf_bert_classification( |
| 214 | num_labels=cfg.num_labels, |
| 215 | pretrained_model_name=cfg.pretrained_model_name, |
| 216 | use_pretrained=cfg.get("use_pretrained", False), |
| 217 | model_config=cfg.get("model_config"), |
| 218 | tokenizer_name=cfg.get("tokenizer_name"), |
| 219 | gradient_checkpointing=cfg.get("gradient_checkpointing"), |
| 220 | ) |
| 221 | elif cfg.name == "mosaic_bert": |
| 222 | return mosaic_bert_module.create_mosaic_bert_classification( |
| 223 | num_labels=cfg.num_labels, |
| 224 | pretrained_model_name=cfg.pretrained_model_name, |
| 225 | pretrained_checkpoint=cfg.get("pretrained_checkpoint"), |
| 226 | model_config=cfg.get("model_config"), |
| 227 | tokenizer_name=cfg.get("tokenizer_name"), |
| 228 | gradient_checkpointing=cfg.get("gradient_checkpointing"), |
| 229 | ) |
| 230 | elif cfg.name == "flex_bert": |
| 231 | return flex_bert_module.create_flex_bert_classification( |
| 232 | num_labels=cfg.num_labels, |
| 233 | pretrained_model_name=cfg.pretrained_model_name, |
| 234 | pretrained_checkpoint=cfg.get("pretrained_checkpoint"), |
| 235 | model_config=cfg.get("model_config"), |
| 236 | tokenizer_name=cfg.get("tokenizer_name"), |
| 237 | gradient_checkpointing=cfg.get("gradient_checkpointing"), |
| 238 | ) |
| 239 | else: |
| 240 | raise ValueError(f"Not sure how to build model with name={cfg.name}") |
| 241 | |
| 242 | |
| 243 | def train(cfg: DictConfig, return_trainer: bool = False, do_train: bool = True) -> Optional[Trainer]: |