| import os | |
| import streamlit as st | |
| from huggingface_hub import HfApi, SpaceHardware | |
| # Set up Hugging Face API token and Space ID | |
| HF_TOKEN = os.getenv("HF_TOKEN") # Ensure your Hugging Face token is set as a secret | |
| TRAINING_SPACE_ID = "your_space_id_here" # Replace with your actual space ID | |
| # Initialize Hugging Face API | |
| api = HfApi(token=HF_TOKEN) | |
| # Function to check for a scheduled task (this is a placeholder for your actual task-checking logic) | |
| def get_task(): | |
| # You can implement logic here to check for scheduled tasks | |
| return None # For example, return None if no task is scheduled | |
| # Function to add a new task (you can implement this depending on your use case) | |
| def add_task(task): | |
| # Logic to add a new task | |
| st.write(f"Task '{task}' added!") | |
| # Function to mark the task as "DONE" (this is a placeholder) | |
| def mark_as_done(task): | |
| # Mark the task as done once it's completed | |
| st.write(f"Task '{task}' completed!") | |
| # Function to simulate training the model (replace with actual training logic) | |
| def train_and_upload(task): | |
| # Implement your model training logic here | |
| st.write(f"Training model with task: {task}") | |
| # Check if there’s an existing task | |
| task = get_task() | |
| if task is None: | |
| # Display Gradio interface to request a new task | |
| def gradio_fn(task): | |
| # On user request, add task and request hardware | |
| add_task(task) | |
| api.request_space_hardware(repo_id=TRAINING_SPACE_ID, hardware=SpaceHardware.T4_MEDIUM) | |
| # Use Streamlit to request a task (Gradio interface or a simple button to simulate this) | |
| task_input = st.text_input("Enter task name", "") | |
| if st.button("Request Task"): | |
| gradio_fn(task_input) | |
| else: | |
| # If a task is available, check for hardware | |
| runtime = api.get_space_runtime(repo_id=TRAINING_SPACE_ID) | |
| if runtime.hardware == SpaceHardware.T4_MEDIUM: | |
| # Fine-tune model on GPU if available | |
| train_and_upload(task) | |
| # Mark task as "DONE" after training | |
| mark_as_done(task) | |
| # Reset to CPU hardware after training | |
| api.request_space_hardware(repo_id=TRAINING_SPACE_ID, hardware=SpaceHardware.CPU_BASIC) | |
| else: | |
| # If GPU hardware is not available, request it | |
| api.request_space_hardware(repo_id=TRAINING_SPACE_ID, hardware=SpaceHardware.T4_MEDIUM) | |