mr-kush/urgency-classifier-retraining

1# train_model.py2from load_dataset import load_dataset_from_hub3from model_pipeline import GrievanceClassifier4from configs import get_config5import time6import os7 8def run_grievance_training_pipeline():9    """10    Load configs, dataset, initialize classifier,11    and run the training pipeline with exception handling.12    Prints status messages for dynamic terminal viewing.13    """14    try:15        print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] Loading configurations...", flush=True)16        configs = get_config()17        # Print a short, non-sensitive summary of configs18        print(f"[{time.strftime('%H:%M:%S')}] Configs loaded: dataset_repo_id={configs.dataset_repo_id}, "19              f"model_checkpoint={configs.model_checkpoint}, hub_model_id={configs.hub_model_id}, "20              f"num_labels={len(configs.label2id)}",21              flush=True)22 23 24        print(f"[{time.strftime('%H:%M:%S')}] Loading dataset from hub: {configs.dataset_repo_id} ...", flush=True)25        data = load_dataset_from_hub(26            model_repo=configs.dataset_repo_id,27            hf_token=configs.hf_token28        )29        dataset = data['dataset']30        dataset_metadata = data['metadata']31 32        # Print dataset splits and sizes if available33        def _safe_len(split):34            try:35                return len(split)36            except Exception:37                return "unknown"38        train_len = _safe_len(dataset.get('train')) if dataset else "no dataset"39        eval_len = _safe_len(dataset.get('eval')) if dataset else "no dataset"40        test_len = _safe_len(dataset.get('test')) if dataset else "no dataset"41        print(f"[{time.strftime('%H:%M:%S')}] Dataset loaded: train={train_len}, eval={eval_len}, test={test_len}", flush=True)42 43        print(f"[{time.strftime('%H:%M:%S')}] Initializing classifier (checkpoint={configs.model_checkpoint}) ...", flush=True)44        classifier = GrievanceClassifier(45            model_checkpoint=configs.model_checkpoint,46            num_labels=len(configs.label2id),47            id2label=configs.id2label,48            label2id=configs.label2id,49            hf_token=configs.hf_token,50            wandb_api_key=configs.wandb_api_key,51            wandb_project_name=configs.wandb_project_name,52        )53        print(f"[{time.strftime('%H:%M:%S')}] Classifier initialized.", flush=True)54 55        print(f"[{time.strftime('%H:%M:%S')}] Start training the model ...", flush=True)56        result = classifier.train_pipeline(57            train_dataset=dataset['train'],58            eval_dataset=dataset['eval'],59            test_dataset=dataset['test'],60            dataset_metadata= dataset_metadata, 61            space_repo_id=configs.space_repo_id,62            hf_training_args={"hub_model_id": configs.hub_model_id},63            api_endpoint=configs.api_endpoint,64            early_stopping_patience=configs.early_stopping_patience,65            deployed_sample_size=configs.deployed_sample_size,66            decision_threshold=configs.decision_threshold67        )68 69        print(f"[{time.strftime('%H:%M:%S')}] Training completed successfully!", flush=True)70        # Print a brief summary of the result if it's a dict-like object71        try:72            if isinstance(result, dict):73                print(f"[{time.strftime('%H:%M:%S')}] Result keys: {list(result.keys())}", flush=True)74            else:75                print(f"[{time.strftime('%H:%M:%S')}] Result: {result}", flush=True)76        except Exception:77            print(f"[{time.strftime('%H:%M:%S')}] Training finished (could not display result details).", flush=True)78 79        # pause the space if it was run in the hf_space80        if configs.retrain_space_id:81            try:82                print(f"[{time.strftime('%H:%M:%S')}] Attempting to pause Hugging Face Space...", flush=True)83                84                classifier.api.pause_space(repo_id=configs.retrain_space_id, token=configs.hf_token)85                86                print(f"[{time.strftime('%H:%M:%S')}] Pause command executed.", flush=True)87            except Exception as e:88                print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] WARNING: Failed to pause HF Space: {e}", flush=True)89 90        return result91 92        93    except Exception as e:94        print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] ERROR: Grievance training pipeline failed: {e}", flush=True)95        raise RuntimeError(f"Grievance training pipeline failed: {e}")96 97 98if __name__ == "__main__":99    run_grievance_training_pipeline()100