Spaces:
Sleeping
Sleeping
| # import os | |
| # import shutil | |
| # import tempfile | |
| # import gradio as gr | |
| # from huggingface_hub import hf_hub_download, upload_file, HfApi | |
| # import sys | |
| # | |
| # # Add current directory to path to import train_model | |
| # sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| # | |
| # # Configuration | |
| # OUTPUT_DIR = "output_data" | |
| # MODEL_FILE = "model_enhanced.pt" | |
| # VOCAB_FILE = "vocabs_enhanced.pkl" | |
| # CHECKPOINT_FILE = "checkpoint_enhanced.pt" | |
| # | |
| # # IMPORTANT: Update this with your actual Hugging Face repository ID | |
| # REPO_ID = "heerjtdev/LSTM_CRF" # Replace with your repo ID | |
| # # HF_TOKEN = os.environ.get("HF_TOKEN") # Set this as a secret in your Space settings | |
| # | |
| # | |
| # def download_existing_models(): | |
| # """Download existing model files from the Hugging Face Hub if available.""" | |
| # try: | |
| # api = HfApi() | |
| # #files = api.list_repo_files(REPO_ID, token=HF_TOKEN) | |
| # files = api.list_repo_files(REPO_ID) | |
| # | |
| # os.makedirs(OUTPUT_DIR, exist_ok=True) | |
| # | |
| # downloaded_files = [] | |
| # | |
| # # Download model file | |
| # if MODEL_FILE in files: | |
| # print(f"📥 Downloading {MODEL_FILE} from Hub...") | |
| # model_path = hf_hub_download( | |
| # repo_id=REPO_ID, | |
| # filename=MODEL_FILE, | |
| # # token=HF_TOKEN, | |
| # local_dir=OUTPUT_DIR, | |
| # force_download=True # Always get latest version | |
| # ) | |
| # downloaded_files.append(MODEL_FILE) | |
| # print(f"✅ Downloaded {MODEL_FILE}") | |
| # | |
| # # Download vocab file | |
| # if VOCAB_FILE in files: | |
| # print(f"📥 Downloading {VOCAB_FILE} from Hub...") | |
| # vocab_path = hf_hub_download( | |
| # repo_id=REPO_ID, | |
| # filename=VOCAB_FILE, | |
| # # token=HF_TOKEN, | |
| # local_dir=OUTPUT_DIR, | |
| # force_download=True # Always get latest version | |
| # ) | |
| # downloaded_files.append(VOCAB_FILE) | |
| # print(f"✅ Downloaded {VOCAB_FILE}") | |
| # | |
| # # Download checkpoint file (optional, for resuming training) | |
| # if CHECKPOINT_FILE in files: | |
| # print(f"📥 Downloading {CHECKPOINT_FILE} from Hub...") | |
| # checkpoint_path = hf_hub_download( | |
| # repo_id=REPO_ID, | |
| # filename=CHECKPOINT_FILE, | |
| # # token=HF_TOKEN, | |
| # local_dir=OUTPUT_DIR, | |
| # force_download=True | |
| # ) | |
| # downloaded_files.append(CHECKPOINT_FILE) | |
| # print(f"✅ Downloaded {CHECKPOINT_FILE}") | |
| # | |
| # if downloaded_files: | |
| # return f"✅ Downloaded from Hub: {', '.join(downloaded_files)}" | |
| # else: | |
| # return "ℹ️ No existing model files found in repository. Starting fresh." | |
| # except Exception as e: | |
| # error_msg = f"⚠️ Could not download existing models: {str(e)}" | |
| # print(error_msg) | |
| # return error_msg | |
| # | |
| # | |
| # def train_model(dataset_file, progress=gr.Progress()): | |
| # """Train the model with the uploaded dataset.""" | |
| # if dataset_file is None: | |
| # return "❌ Please upload a dataset file!", None, None | |
| # | |
| # try: | |
| # # Step 1: Download existing models from Hub (if any) BEFORE training starts | |
| # progress(0.05, desc="Checking Hugging Face Hub for existing models...") | |
| # download_status = download_existing_models() | |
| # status_log = f"{download_status}\n\n" | |
| # yield status_log, None, None | |
| # | |
| # # Step 2: Save uploaded file | |
| # progress(0.1, desc="Processing uploaded dataset...") | |
| # dataset_path = dataset_file.name | |
| # status_log += f"📂 Dataset uploaded: {os.path.basename(dataset_path)}\n\n" | |
| # yield status_log, None, None | |
| # | |
| # # Step 3: Import and run training | |
| # progress(0.15, desc="Initializing training...") | |
| # status_log += "🚀 Starting training...\n" | |
| # status_log += "📊 This may take a while. Training progress will appear in the terminal.\n\n" | |
| # yield status_log, None, None | |
| # | |
| # # Import the training module | |
| # try: | |
| # import train_model as tm | |
| # print("=" * 80) | |
| # print("TRAINING STARTED") | |
| # print("=" * 80) | |
| # | |
| # # Run training - this will handle model loading internally | |
| # progress(0.2, desc="Training in progress... (check terminal for details)") | |
| # tm.train_from_json(dataset_path) | |
| # | |
| # print("=" * 80) | |
| # print("TRAINING COMPLETED") | |
| # print("=" * 80) | |
| # | |
| # status_log += "✅ Training completed successfully!\n\n" | |
| # yield status_log, None, None | |
| # | |
| # except ImportError as ie: | |
| # error_msg = f"❌ Failed to import training module: {str(ie)}\n" | |
| # error_msg += "Make sure train_model.py is in the same directory as app.py" | |
| # yield status_log + error_msg, None, None | |
| # return | |
| # except Exception as train_error: | |
| # error_msg = f"❌ Training failed with error:\n{str(train_error)}\n" | |
| # yield status_log + error_msg, None, None | |
| # return | |
| # | |
| # # Step 4: Verify files exist | |
| # progress(0.85, desc="Verifying trained model files...") | |
| # model_path = os.path.join(OUTPUT_DIR, MODEL_FILE) | |
| # vocab_path = os.path.join(OUTPUT_DIR, VOCAB_FILE) | |
| # checkpoint_path = os.path.join(OUTPUT_DIR, CHECKPOINT_FILE) | |
| # | |
| # files_exist = [] | |
| # if os.path.exists(model_path): | |
| # files_exist.append(MODEL_FILE) | |
| # if os.path.exists(vocab_path): | |
| # files_exist.append(VOCAB_FILE) | |
| # | |
| # if not files_exist: | |
| # error_msg = "❌ Error: Model files were not created. Check training logs." | |
| # yield status_log + error_msg, None, None | |
| # return | |
| # | |
| # status_log += f"✅ Found trained files: {', '.join(files_exist)}\n\n" | |
| # yield status_log, None, None | |
| # | |
| # # Step 5: Upload to Hub | |
| # progress(0.9, desc="Uploading models to Hugging Face Hub...") | |
| # status_log += "☁️ Uploading to Hugging Face Hub...\n" | |
| # yield status_log, None, None | |
| # | |
| # upload_status = [] | |
| # | |
| # if os.path.exists(model_path): | |
| # try: | |
| # upload_file( | |
| # path_or_fileobj=model_path, | |
| # path_in_repo=MODEL_FILE, | |
| # repo_id=REPO_ID, | |
| # # token=HF_TOKEN, | |
| # commit_message="Update trained model" | |
| # ) | |
| # upload_status.append(MODEL_FILE) | |
| # print(f"✅ Uploaded {MODEL_FILE} to Hub") | |
| # except Exception as e: | |
| # print(f"⚠️ Failed to upload {MODEL_FILE}: {e}") | |
| # | |
| # if os.path.exists(vocab_path): | |
| # try: | |
| # upload_file( | |
| # path_or_fileobj=vocab_path, | |
| # path_in_repo=VOCAB_FILE, | |
| # repo_id=REPO_ID, | |
| # # token=HF_TOKEN, | |
| # commit_message="Update vocabulary" | |
| # ) | |
| # upload_status.append(VOCAB_FILE) | |
| # print(f"✅ Uploaded {VOCAB_FILE} to Hub") | |
| # except Exception as e: | |
| # print(f"⚠️ Failed to upload {VOCAB_FILE}: {e}") | |
| # | |
| # # Also upload checkpoint for future resume capability | |
| # if os.path.exists(checkpoint_path): | |
| # try: | |
| # upload_file( | |
| # path_or_fileobj=checkpoint_path, | |
| # path_in_repo=CHECKPOINT_FILE, | |
| # repo_id=REPO_ID, | |
| # # token=HF_TOKEN, | |
| # commit_message="Update checkpoint" | |
| # ) | |
| # upload_status.append(CHECKPOINT_FILE) | |
| # print(f"✅ Uploaded {CHECKPOINT_FILE} to Hub") | |
| # except Exception as e: | |
| # print(f"⚠️ Failed to upload {CHECKPOINT_FILE}: {e}") | |
| # | |
| # if upload_status: | |
| # status_log += f"✅ Uploaded to Hub: {', '.join(upload_status)}\n\n" | |
| # else: | |
| # status_log += "⚠️ Warning: No files were uploaded to Hub\n\n" | |
| # | |
| # yield status_log, None, None | |
| # | |
| # # Step 6: Copy to temp directory for download | |
| # progress(0.95, desc="Preparing download files...") | |
| # temp_dir = tempfile.mkdtemp() | |
| # | |
| # model_download = None | |
| # vocab_download = None | |
| # | |
| # if os.path.exists(model_path): | |
| # temp_model = os.path.join(temp_dir, MODEL_FILE) | |
| # shutil.copy2(model_path, temp_model) | |
| # model_download = temp_model | |
| # print(f"📦 Prepared {MODEL_FILE} for download") | |
| # | |
| # if os.path.exists(vocab_path): | |
| # temp_vocab = os.path.join(temp_dir, VOCAB_FILE) | |
| # shutil.copy2(vocab_path, temp_vocab) | |
| # vocab_download = temp_vocab | |
| # print(f"📦 Prepared {VOCAB_FILE} for download") | |
| # | |
| # progress(1.0, desc="Complete!") | |
| # | |
| # status_log += "📦 Files ready for download below!\n" | |
| # status_log += "\n" + "=" * 50 + "\n" | |
| # status_log += "TRAINING COMPLETE - You can now download the model files\n" | |
| # status_log += "=" * 50 | |
| # | |
| # yield status_log, model_download, vocab_download | |
| # | |
| # except Exception as e: | |
| # error_msg = f"❌ Unexpected error: {str(e)}\n" | |
| # import traceback | |
| # error_msg += f"\nTraceback:\n{traceback.format_exc()}" | |
| # yield error_msg, None, None | |
| # | |
| # | |
| # def download_models_from_hub(): | |
| # """Download the latest models from the Hugging Face Hub.""" | |
| # try: | |
| # os.makedirs(OUTPUT_DIR, exist_ok=True) | |
| # | |
| # api = HfApi() | |
| # #files = api.list_repo_files(REPO_ID, token=HF_TOKEN) | |
| # files = api.list_repo_files(REPO_ID) | |
| # | |
| # downloaded_files = [] | |
| # | |
| # # Download model | |
| # if MODEL_FILE in files: | |
| # print(f"📥 Downloading {MODEL_FILE} from Hub...") | |
| # model_path = hf_hub_download( | |
| # repo_id=REPO_ID, | |
| # filename=MODEL_FILE, | |
| # # token=HF_TOKEN, | |
| # local_dir=OUTPUT_DIR, | |
| # force_download=True | |
| # ) | |
| # downloaded_files.append(MODEL_FILE) | |
| # else: | |
| # return f"❌ {MODEL_FILE} not found in repository", None, None | |
| # | |
| # # Download vocab | |
| # if VOCAB_FILE in files: | |
| # print(f"📥 Downloading {VOCAB_FILE} from Hub...") | |
| # vocab_path = hf_hub_download( | |
| # repo_id=REPO_ID, | |
| # filename=VOCAB_FILE, | |
| # # token=HF_TOKEN, | |
| # local_dir=OUTPUT_DIR, | |
| # force_download=True | |
| # ) | |
| # downloaded_files.append(VOCAB_FILE) | |
| # else: | |
| # return f"❌ {VOCAB_FILE} not found in repository", None, None | |
| # | |
| # # Copy to temp for download | |
| # temp_dir = tempfile.mkdtemp() | |
| # temp_model = os.path.join(temp_dir, MODEL_FILE) | |
| # temp_vocab = os.path.join(temp_dir, VOCAB_FILE) | |
| # | |
| # shutil.copy2(os.path.join(OUTPUT_DIR, MODEL_FILE), temp_model) | |
| # shutil.copy2(os.path.join(OUTPUT_DIR, VOCAB_FILE), temp_vocab) | |
| # | |
| # success_msg = f"✅ Successfully downloaded from Hub:\n" | |
| # success_msg += f" • {MODEL_FILE}\n" | |
| # success_msg += f" • {VOCAB_FILE}\n\n" | |
| # success_msg += "📦 Files are ready to download below!" | |
| # | |
| # return success_msg, temp_model, temp_vocab | |
| # | |
| # except Exception as e: | |
| # error_msg = f"❌ Error downloading models: {str(e)}\n\n" | |
| # error_msg += f"Make sure:\n" | |
| # error_msg += f"1. REPO_ID is set correctly: {REPO_ID}\n" | |
| # error_msg += f"2. HF_TOKEN is set in Space secrets\n" | |
| # error_msg += f"3. Model files exist in the repository" | |
| # return error_msg, None, None | |
| # | |
| # | |
| # # Create Gradio interface | |
| # with gr.Blocks(title="MCQ Structure Extraction - Model Training", theme=gr.themes.Soft()) as demo: | |
| # gr.Markdown( | |
| # """ | |
| # # 🎓 MCQ Structure Extraction - Model Training | |
| # | |
| # Train a BiLSTM-CRF model with deep layout understanding for extracting structured information from MCQ documents. | |
| # | |
| # ## 📋 Instructions: | |
| # 1. **Upload Dataset**: Provide your unified JSON file containing tokens, bounding boxes, and labels | |
| # 2. **Train Model**: Click "Start Training" and wait for completion (this may take a while) | |
| # 3. **Download Models**: Once training is complete, download the trained model and vocabulary files | |
| # | |
| # ## 📥 Or Download Existing Models: | |
| # If you just want to download the latest trained models from the repository, use the "Download from Hub" tab. | |
| # | |
| # --- | |
| # """ | |
| # ) | |
| # | |
| # with gr.Tab("🚀 Train New Model"): | |
| # gr.Markdown( | |
| # """ | |
| # ### Training Process: | |
| # The app will automatically: | |
| # 1. ✅ Download any existing models from Hugging Face Hub (for resuming training) | |
| # 2. 🎯 Train the model on your uploaded dataset | |
| # 3. ☁️ Upload the trained models back to the Hub | |
| # 4. 📥 Provide download links for the trained files | |
| # | |
| # **Note**: Training progress details appear in the terminal/logs. The status box shows major milestones. | |
| # """ | |
| # ) | |
| # | |
| # with gr.Row(): | |
| # with gr.Column(): | |
| # dataset_input = gr.File( | |
| # label="📂 Upload Training Dataset (JSON)", | |
| # file_types=[".json"], | |
| # type="filepath" | |
| # ) | |
| # train_button = gr.Button("🚀 Start Training", variant="primary", size="lg") | |
| # | |
| # with gr.Column(): | |
| # status_output = gr.Textbox( | |
| # label="📊 Training Status", | |
| # lines=12, | |
| # interactive=False, | |
| # show_copy_button=True | |
| # ) | |
| # | |
| # gr.Markdown("### 📦 Download Trained Models") | |
| # with gr.Row(): | |
| # model_output = gr.File(label="💾 Model File (.pt)") | |
| # vocab_output = gr.File(label="📚 Vocabulary File (.pkl)") | |
| # | |
| # train_button.click( | |
| # fn=train_model, | |
| # inputs=[dataset_input], | |
| # outputs=[status_output, model_output, vocab_output] | |
| # ) | |
| # | |
| # with gr.Tab("☁️ Download from Hub"): | |
| # gr.Markdown( | |
| # """ | |
| # ### Download Pre-trained Models | |
| # | |
| # Download the latest trained models directly from your Hugging Face repository. | |
| # This is useful if: | |
| # - You want to use pre-trained models without training | |
| # - You need to download models trained in a previous session | |
| # - You want to get the latest version from the Hub | |
| # | |
| # The downloaded files can be used for inference with your MCQ extraction pipeline. | |
| # """ | |
| # ) | |
| # | |
| # download_button = gr.Button("☁️ Download Latest Models from Hub", variant="primary", size="lg") | |
| # | |
| # download_status = gr.Textbox( | |
| # label="Download Status", | |
| # lines=6, | |
| # interactive=False, | |
| # show_copy_button=True | |
| # ) | |
| # | |
| # gr.Markdown("### 📦 Downloaded Files") | |
| # with gr.Row(): | |
| # hub_model_output = gr.File(label="💾 Model File (.pt)") | |
| # hub_vocab_output = gr.File(label="📚 Vocabulary File (.pkl)") | |
| # | |
| # download_button.click( | |
| # fn=download_models_from_hub, | |
| # outputs=[download_status, hub_model_output, hub_vocab_output] | |
| # ) | |
| # | |
| # gr.Markdown( | |
| # """ | |
| # --- | |
| # ### ⚙️ Model Configuration: | |
| # | |
| # **Architecture:** | |
| # - BiLSTM-CRF with spatial attention mechanism | |
| # - Word embeddings + Character-level CNN | |
| # - Bounding box encoding with MLP | |
| # - Spatial & context feature extraction | |
| # - Learnable positional embeddings | |
| # | |
| # **Features Used:** | |
| # - Token text (word-level and character-level) | |
| # - Bounding box coordinates (normalized) | |
| # - Spatial features: vertical spacing, alignment, dimensions (11 features) | |
| # - Context features: surrounding question/option markers (8 features) | |
| # | |
| # **Output Labels (13 total):** | |
| # - Questions, Options, Answers, Images, Section Headings, Passages (BIO tagging) | |
| # | |
| # **Training Parameters:** | |
| # - Batch Size: 8 | |
| # - Epochs: 10 (with early stopping after 10 epochs without improvement) | |
| # - Learning Rate: 5e-4 (AdamW optimizer with OneCycleLR scheduler) | |
| # - Hidden Size: 768 | |
| # - Total Parameters: ~15.6M | |
| # | |
| # **Hardware Requirements:** | |
| # - GPU recommended for reasonable training speed | |
| # - CPU training supported but significantly slower | |
| # | |
| # --- | |
| # | |
| # | |
| # | |
| # **Environment Variables Required:** | |
| # - `SPACE_ID`: Your Hugging Face Space/Repo ID (auto-set in Spaces) | |
| # - `HF_TOKEN`: Your Hugging Face write token (set as a secret) | |
| # | |
| # **Model Persistence:** | |
| # - Models are automatically saved to `output_data/` directory | |
| # - Best model is uploaded to Hugging Face Hub after each improvement | |
| # - Training can be resumed from checkpoints | |
| # """ | |
| # ) | |
| # | |
| # # Launch the app | |
| # if __name__ == "__main__": | |
| # demo.launch() | |
| import os | |
| import shutil | |
| import tempfile | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download, upload_file, HfApi | |
| import sys | |
| import glob | |
| # Add current directory to path to import train_model | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| # Configuration | |
| OUTPUT_DIR = "output_data" | |
| MODEL_FILE = "model_enhanced.pt" | |
| VOCAB_FILE = "vocabs_enhanced.pkl" | |
| CHECKPOINT_FILE = "checkpoint_enhanced.pt" | |
| # IMPORTANT: Update this with your actual Hugging Face repository ID | |
| REPO_ID = "heerjtdev/LSTM_CRF" # Replace with your repo ID | |
| # HF_TOKEN = os.environ.get("HF_TOKEN") # Set this as a secret in your Space settings | |
| def download_existing_models(): | |
| """Download existing model files from the Hugging Face Hub if available.""" | |
| try: | |
| api = HfApi() | |
| # files = api.list_repo_files(REPO_ID, token=HF_TOKEN) | |
| files = api.list_repo_files(REPO_ID) | |
| os.makedirs(OUTPUT_DIR, exist_ok=True) | |
| downloaded_files = [] | |
| # Download model file | |
| if MODEL_FILE in files: | |
| print(f"📥 Downloading {MODEL_FILE} from Hub...") | |
| model_path = hf_hub_download( | |
| repo_id=REPO_ID, | |
| filename=MODEL_FILE, | |
| # token=HF_TOKEN, | |
| local_dir=OUTPUT_DIR, | |
| force_download=True # Always get latest version | |
| ) | |
| downloaded_files.append(MODEL_FILE) | |
| print(f"✅ Downloaded {MODEL_FILE}") | |
| # Download vocab file | |
| if VOCAB_FILE in files: | |
| print(f"📥 Downloading {VOCAB_FILE} from Hub...") | |
| vocab_path = hf_hub_download( | |
| repo_id=REPO_ID, | |
| filename=VOCAB_FILE, | |
| # token=HF_TOKEN, | |
| local_dir=OUTPUT_DIR, | |
| force_download=True # Always get latest version | |
| ) | |
| downloaded_files.append(VOCAB_FILE) | |
| print(f"✅ Downloaded {VOCAB_FILE}") | |
| # Download checkpoint file (optional, for resuming training) | |
| if CHECKPOINT_FILE in files: | |
| print(f"📥 Downloading {CHECKPOINT_FILE} from Hub...") | |
| checkpoint_path = hf_hub_download( | |
| repo_id=REPO_ID, | |
| filename=CHECKPOINT_FILE, | |
| # token=HF_TOKEN, | |
| local_dir=OUTPUT_DIR, | |
| force_download=True | |
| ) | |
| downloaded_files.append(CHECKPOINT_FILE) | |
| print(f"✅ Downloaded {CHECKPOINT_FILE}") | |
| if downloaded_files: | |
| return f"✅ Downloaded from Hub: {', '.join(downloaded_files)}" | |
| else: | |
| return "ℹ️ No existing model files found in repository. Starting fresh." | |
| except Exception as e: | |
| error_msg = f"⚠️ Could not download existing models: {str(e)}" | |
| print(error_msg) | |
| return error_msg | |
| def train_model(dataset_file, progress=gr.Progress()): | |
| """Train the model with the uploaded dataset.""" | |
| if dataset_file is None: | |
| return "❌ Please upload a dataset file!", None, None | |
| try: | |
| # Step 1: Download existing models from Hub (if any) BEFORE training starts | |
| progress(0.05, desc="Checking Hugging Face Hub for existing models...") | |
| download_status = download_existing_models() | |
| status_log = f"{download_status}\n\n" | |
| # Reset download outputs before training starts | |
| yield status_log, None, None, None, None | |
| # Step 2: Save uploaded file | |
| progress(0.1, desc="Processing uploaded dataset...") | |
| dataset_path = dataset_file.name | |
| status_log += f"📂 Dataset uploaded: {os.path.basename(dataset_path)}\n\n" | |
| yield status_log, None, None, None, None | |
| # Step 3: Import and run training | |
| progress(0.15, desc="Initializing training...") | |
| status_log += "🚀 Starting training...\n" | |
| status_log += "📊 This may take a while. Training progress will appear in the terminal.\n\n" | |
| yield status_log, None, None, None, None | |
| # Import the training module | |
| try: | |
| import train_model as tm | |
| print("=" * 80) | |
| print("TRAINING STARTED") | |
| print("=" * 80) | |
| # Run training - this will handle model loading internally | |
| progress(0.2, desc="Training in progress... (check terminal for details)") | |
| tm.train_from_json(dataset_path) | |
| print("=" * 80) | |
| print("TRAINING COMPLETED") | |
| print("=" * 80) | |
| status_log += "✅ Training completed successfully!\n\n" | |
| yield status_log, None, None, None, None | |
| except ImportError as ie: | |
| error_msg = f"❌ Failed to import training module: {str(ie)}\n" | |
| error_msg += "Make sure train_model.py is in the same directory as app.py" | |
| yield status_log + error_msg, None, None, None, None | |
| return | |
| except Exception as train_error: | |
| error_msg = f"❌ Training failed with error:\n{str(train_error)}\n" | |
| yield status_log + error_msg, None, None, None, None | |
| return | |
| # Step 4: Verify files exist | |
| progress(0.85, desc="Verifying trained model files...") | |
| model_path = os.path.join(OUTPUT_DIR, MODEL_FILE) | |
| vocab_path = os.path.join(OUTPUT_DIR, VOCAB_FILE) | |
| checkpoint_path = os.path.join(OUTPUT_DIR, CHECKPOINT_FILE) | |
| files_exist = [] | |
| if os.path.exists(model_path): | |
| files_exist.append(MODEL_FILE) | |
| if os.path.exists(vocab_path): | |
| files_exist.append(VOCAB_FILE) | |
| if not files_exist: | |
| error_msg = "❌ Error: Model files were not created. Check training logs." | |
| yield status_log + error_msg, None, None, None, None | |
| return | |
| status_log += f"✅ Found trained files: {', '.join(files_exist)}\n\n" | |
| yield status_log, None, None, None, None | |
| # Step 5: Upload to Hub | |
| progress(0.9, desc="Uploading models to Hugging Face Hub...") | |
| status_log += "☁️ Uploading to Hugging Face Hub...\n" | |
| yield status_log, None, None, None, None | |
| upload_status = [] | |
| if os.path.exists(model_path): | |
| try: | |
| upload_file( | |
| path_or_fileobj=model_path, | |
| path_in_repo=MODEL_FILE, | |
| repo_id=REPO_ID, | |
| # token=HF_TOKEN, | |
| commit_message="Update trained model" | |
| ) | |
| upload_status.append(MODEL_FILE) | |
| print(f"✅ Uploaded {MODEL_FILE} to Hub") | |
| except Exception as e: | |
| print(f"⚠️ Failed to upload {MODEL_FILE}: {e}") | |
| if os.path.exists(vocab_path): | |
| try: | |
| upload_file( | |
| path_or_fileobj=vocab_path, | |
| path_in_repo=VOCAB_FILE, | |
| repo_id=REPO_ID, | |
| # token=HF_TOKEN, | |
| commit_message="Update vocabulary" | |
| ) | |
| upload_status.append(VOCAB_FILE) | |
| print(f"✅ Uploaded {VOCAB_FILE} to Hub") | |
| except Exception as e: | |
| print(f"⚠️ Failed to upload {VOCAB_FILE}: {e}") | |
| # Also upload checkpoint for future resume capability | |
| if os.path.exists(checkpoint_path): | |
| try: | |
| upload_file( | |
| path_or_fileobj=checkpoint_path, | |
| path_in_repo=CHECKPOINT_FILE, | |
| repo_id=REPO_ID, | |
| # token=HF_TOKEN, | |
| commit_message="Update checkpoint" | |
| ) | |
| upload_status.append(CHECKPOINT_FILE) | |
| print(f"✅ Uploaded {CHECKPOINT_FILE} to Hub") | |
| except Exception as e: | |
| print(f"⚠️ Failed to upload {CHECKPOINT_FILE}: {e}") | |
| if upload_status: | |
| status_log += f"✅ Uploaded to Hub: {', '.join(upload_status)}\n\n" | |
| else: | |
| status_log += "⚠️ Warning: No files were uploaded to Hub\n\n" | |
| yield status_log, None, None, None, None | |
| # Step 6: Copy to temp directory for download | |
| progress(0.95, desc="Preparing download files...") | |
| temp_dir = tempfile.mkdtemp() | |
| model_download = None | |
| vocab_download = None | |
| if os.path.exists(model_path): | |
| temp_model = os.path.join(temp_dir, MODEL_FILE) | |
| shutil.copy2(model_path, temp_model) | |
| model_download = temp_model | |
| print(f"📦 Prepared {MODEL_FILE} for download") | |
| if os.path.exists(vocab_path): | |
| temp_vocab = os.path.join(temp_dir, VOCAB_FILE) | |
| shutil.copy2(vocab_path, temp_vocab) | |
| vocab_download = temp_vocab | |
| print(f"📦 Prepared {VOCAB_FILE} for download") | |
| progress(1.0, desc="Complete!") | |
| status_log += "📦 Files ready for download below!\n" | |
| status_log += "\n" + "=" * 50 + "\n" | |
| status_log += "TRAINING COMPLETE - You can now download the model files\n" | |
| status_log += "=" * 50 | |
| # Note: We return the model_download and vocab_download twice for both sets of File outputs | |
| yield status_log, model_download, vocab_download, model_download, vocab_download | |
| except Exception as e: | |
| error_msg = f"❌ Unexpected error: {str(e)}\n" | |
| import traceback | |
| error_msg += f"\nTraceback:\n{traceback.format_exc()}" | |
| # Return Nones for all file outputs | |
| yield error_msg, None, None, None, None | |
| def download_models_from_hub(): | |
| """Download the latest models from the Hugging Face Hub.""" | |
| try: | |
| os.makedirs(OUTPUT_DIR, exist_ok=True) | |
| api = HfApi() | |
| # files = api.list_repo_files(REPO_ID, token=HF_TOKEN) | |
| files = api.list_repo_files(REPO_ID) | |
| downloaded_files = [] | |
| # Download model | |
| if MODEL_FILE in files: | |
| print(f"📥 Downloading {MODEL_FILE} from Hub...") | |
| model_path = hf_hub_download( | |
| repo_id=REPO_ID, | |
| filename=MODEL_FILE, | |
| # token=HF_TOKEN, | |
| local_dir=OUTPUT_DIR, | |
| force_download=True | |
| ) | |
| downloaded_files.append(MODEL_FILE) | |
| else: | |
| return f"❌ {MODEL_FILE} not found in repository", None, None, None, None | |
| # Download vocab | |
| if VOCAB_FILE in files: | |
| print(f"📥 Downloading {VOCAB_FILE} from Hub...") | |
| vocab_path = hf_hub_download( | |
| repo_id=REPO_ID, | |
| filename=VOCAB_FILE, | |
| # token=HF_TOKEN, | |
| local_dir=OUTPUT_DIR, | |
| force_download=True | |
| ) | |
| downloaded_files.append(VOCAB_FILE) | |
| else: | |
| return f"❌ {VOCAB_FILE} not found in repository", None, None, None, None | |
| # Copy to temp for download | |
| temp_dir = tempfile.mkdtemp() | |
| temp_model = os.path.join(temp_dir, MODEL_FILE) | |
| temp_vocab = os.path.join(temp_dir, VOCAB_FILE) | |
| shutil.copy2(os.path.join(OUTPUT_DIR, MODEL_FILE), temp_model) | |
| shutil.copy2(os.path.join(OUTPUT_DIR, VOCAB_FILE), temp_vocab) | |
| success_msg = f"✅ Successfully downloaded from Hub:\n" | |
| success_msg += f" • {MODEL_FILE}\n" | |
| success_msg += f" • {VOCAB_FILE}\n\n" | |
| success_msg += "📦 Files are ready to download below!" | |
| # Return the downloaded files for both sets of file outputs | |
| return success_msg, temp_model, temp_vocab, temp_model, temp_vocab | |
| except Exception as e: | |
| error_msg = f"❌ Error downloading models: {str(e)}\n\n" | |
| error_msg += f"Make sure:\n" | |
| error_msg += f"1. REPO_ID is set correctly: {REPO_ID}\n" | |
| error_msg += f"2. HF_TOKEN is set in Space secrets\n" | |
| error_msg += f"3. Model files exist in the repository" | |
| return error_msg, None, None, None, None | |
| # --- UPDATED check_local_files FUNCTION --- | |
| def check_local_files(): | |
| """ | |
| Checks and reports the files present in the local output directory. | |
| If core model files exist, it prepares and returns them for download. | |
| """ | |
| if not os.path.exists(OUTPUT_DIR): | |
| return f"ℹ️ Directory **'{OUTPUT_DIR}'** does not exist.", None, None | |
| all_files = os.listdir(OUTPUT_DIR) | |
| model_path = os.path.join(OUTPUT_DIR, MODEL_FILE) | |
| vocab_path = os.path.join(OUTPUT_DIR, VOCAB_FILE) | |
| model_download = None | |
| vocab_download = None | |
| # 1. Prepare download paths if files exist | |
| if os.path.exists(model_path): | |
| model_download = model_path | |
| if os.path.exists(vocab_path): | |
| vocab_download = vocab_path | |
| # 2. Generate status message | |
| if not all_files: | |
| return f"ℹ️ Directory **'{OUTPUT_DIR}'** is empty.", None, None | |
| file_list = [] | |
| total_size = 0 | |
| # Sort files to put core files first | |
| sorted_files = sorted(all_files, key=lambda x: (x != MODEL_FILE, x != VOCAB_FILE, x != CHECKPOINT_FILE, x)) | |
| for filename in sorted_files: | |
| filepath = os.path.join(OUTPUT_DIR, filename) | |
| if os.path.isfile(filepath): | |
| size_bytes = os.path.getsize(filepath) | |
| total_size += size_bytes | |
| # Simple size formatting | |
| if size_bytes > 1024 * 1024: | |
| size_str = f"{size_bytes / (1024 * 1024):.2f} MB" | |
| elif size_bytes > 1024: | |
| size_str = f"{size_bytes / 1024:.2f} KB" | |
| else: | |
| size_str = f"{size_bytes} bytes" | |
| file_list.append(f"• **{filename}** (Size: {size_str})") | |
| # Format total size | |
| if total_size > 1024 * 1024 * 1024: | |
| total_size_str = f"{total_size / (1024 * 1024 * 1024):.2f} GB" | |
| elif total_size > 1024 * 1024: | |
| total_size_str = f"{total_size / (1024 * 1024):.2f} MB" | |
| else: | |
| total_size_str = f"{total_size / 1024:.2f} KB" | |
| header = f"✅ Contents of **'{OUTPUT_DIR}'** ({len(file_list)} files, Total Size: {total_size_str}):\n" | |
| if model_download and vocab_download: | |
| header += "\n**📦 Core model files found! Ready for download below.**" | |
| elif model_download or vocab_download: | |
| header += "\n**⚠️ Found some model files, but not both.**" | |
| return header + "\n" + "\n".join(file_list), model_download, vocab_download | |
| def clear_local_memory(): | |
| """Deletes the local output directory and its contents.""" | |
| if os.path.exists(OUTPUT_DIR): | |
| try: | |
| shutil.rmtree(OUTPUT_DIR) | |
| return f"🗑️ Successfully deleted local directory **'{OUTPUT_DIR}'** and all its contents. Memory cleared.", None, None | |
| except Exception as e: | |
| return f"❌ Error clearing memory (deleting '{OUTPUT_DIR}'): {str(e)}", None, None | |
| else: | |
| return f"ℹ️ Local directory **'{OUTPUT_DIR}'** does not exist. No memory to clear.", None, None | |
| # --- END NEW FUNCTIONS --- | |
| # Create Gradio interface | |
| with gr.Blocks(title="MCQ Structure Extraction - Model Training", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown( | |
| """ | |
| # 🎓 MCQ Structure Extraction - Model Training | |
| Train a BiLSTM-CRF model with deep layout understanding for extracting structured information from MCQ documents. | |
| ## 📋 Instructions: | |
| 1. **Upload Dataset**: Provide your unified JSON file containing tokens, bounding boxes, and labels | |
| 2. **Train Model**: Click "Start Training" and wait for completion (this may take a while) | |
| 3. **Download Models**: Once training is complete, download the trained model and vocabulary files | |
| ## 📥 Or Download Existing Models: | |
| If you just want to download the latest trained models from the repository, use the "Download from Hub" tab. | |
| --- | |
| """ | |
| ) | |
| # Define common File components for outputs | |
| download_model_output = gr.File(label="💾 Model File (.pt)", interactive=False) | |
| download_vocab_output = gr.File(label="📚 Vocabulary File (.pkl)", interactive=False) | |
| # We need a dummy set of outputs to clear the download boxes when starting training, | |
| # and a permanent set for the utility functions. We'll use the permanent ones below. | |
| with gr.Tab("🚀 Train New Model"): | |
| gr.Markdown( | |
| """ | |
| ### Training Process: | |
| The app will automatically: | |
| 1. ✅ Download any existing models from Hugging Face Hub (for resuming training) | |
| 2. 🎯 Train the model on your uploaded dataset | |
| 3. ☁️ Upload the trained models back to the Hub | |
| 4. 📥 Provide download links for the trained files | |
| **Note**: Training progress details appear in the terminal/logs. The status box shows major milestones. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| dataset_input = gr.File( | |
| label="📂 Upload Training Dataset (JSON)", | |
| file_types=[".json"], | |
| type="filepath" | |
| ) | |
| train_button = gr.Button("🚀 Start Training", variant="primary", size="lg") | |
| # --- NEW BUTTONS for utility --- | |
| with gr.Row(): | |
| check_button = gr.Button("🔎 Check Local Models", variant="secondary") | |
| clear_button = gr.Button("🧹 Clear Local Memory", variant="stop") | |
| # ------------------------------ | |
| with gr.Column(): | |
| status_output = gr.Textbox( | |
| label="📊 Training/Utility Status", | |
| lines=12, | |
| interactive=False, | |
| show_copy_button=True | |
| ) | |
| gr.Markdown("### 📦 Download Trained/Local Models") | |
| with gr.Row(): | |
| # Use the defined components for the training output | |
| train_model_output = download_model_output | |
| train_vocab_output = download_vocab_output | |
| # Note: The train_model function now returns 5 values (status, model_file, vocab_file, model_file_again, vocab_file_again) | |
| # We target the two download outputs directly for the final model and vocab files. | |
| train_button.click( | |
| fn=train_model, | |
| inputs=[dataset_input], | |
| outputs=[status_output, train_model_output, train_vocab_output, download_model_output, | |
| download_vocab_output] | |
| ) | |
| # --- NEW BUTTON ACTIONS --- | |
| # check_local_files now returns (status, model_download_path, vocab_download_path) | |
| # We target the status output AND the two global download outputs | |
| check_button.click( | |
| fn=check_local_files, | |
| inputs=[], | |
| outputs=[status_output, download_model_output, download_vocab_output] | |
| ) | |
| # clear_local_memory now returns (status, None, None) to clear the download boxes | |
| clear_button.click( | |
| fn=clear_local_memory, | |
| inputs=[], | |
| outputs=[status_output, download_model_output, download_vocab_output] | |
| ) | |
| # -------------------------- | |
| with gr.Tab("☁️ Download from Hub"): | |
| gr.Markdown( | |
| """ | |
| ### Download Pre-trained Models | |
| Download the latest trained models directly from your Hugging Face repository. | |
| """ | |
| ) | |
| download_button = gr.Button("☁️ Download Latest Models from Hub", variant="primary", size="lg") | |
| download_status = gr.Textbox( | |
| label="Download Status", | |
| lines=6, | |
| interactive=False, | |
| show_copy_button=True | |
| ) | |
| gr.Markdown("### 📦 Downloaded Files") | |
| with gr.Row(): | |
| # Use the defined components for the Hub output | |
| hub_model_output = download_model_output | |
| hub_vocab_output = download_vocab_output | |
| # Note: The download_models_from_hub function now returns 5 values (status, model_file, vocab_file, model_file_again, vocab_file_again) | |
| # We target the two download outputs directly for the final model and vocab files. | |
| download_button.click( | |
| fn=download_models_from_hub, | |
| outputs=[download_status, hub_model_output, hub_vocab_output, download_model_output, download_vocab_output] | |
| ) | |
| gr.Markdown( | |
| """ | |
| --- | |
| ### ⚙️ Model Configuration: | |
| ... (rest of the markdown) | |
| """ | |
| ) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| demo.launch() |