Backend refactor (#627)

configs/dreambooth/sd15_colab.yml

@@ -1,7 +1,7 @@
 task: dreambooth
 base_model: runwayml/stable-diffusion-v1-5
 project_name: autotrain-sd15-finetuned
-backend: local-cli
+backend: local
 
 data:
   path: data/ # store all images in this folder

configs/dreambooth/sdxl_colab.yml

@@ -1,7 +1,7 @@
 task: dreambooth
 base_model: stabilityai/stable-diffusion-xl-base-1.0
 project_name: autotrain-sdxl-finetuned
-backend: local-cli
+backend: local
 
 data:
   path: data/ # store all images in this folder

configs/image_classification/hub_dataset.yml

@@ -2,7 +2,7 @@ task: image_classification
 base_model: google/vit-base-patch16-224
 project_name: autotrain-cats-vs-dogs-finetuned
 log: tensorboard
-backend: local-cli
+backend: local
 
 data:
   path: cats_vs_dogs

configs/llm_finetuning/gpt2_sft.yml

@@ -2,7 +2,7 @@ task: llm
 base_model: openai-community/gpt2
 project_name: autotrain-gpt2-finetuned-guanaco
 log: tensorboard
-backend: local-cli
+backend: local
 
 data:
   path: timdettmers/openassistant-guanaco

configs/llm_finetuning/llama3-70b-orpo-v1.yml

@@ -2,7 +2,7 @@ task: llm
 base_model: meta-llama/Meta-Llama-3-70B-Instruct
 project_name: autotrain-llama3-70b-orpo-v1
 log: tensorboard
-backend: local-cli
+backend: local
 
 data:
   path: argilla/distilabel-capybara-dpo-7k-binarized

configs/llm_finetuning/llama3-70b-sft.yml

@@ -2,7 +2,7 @@ task: llm
 base_model: meta-llama/Meta-Llama-3-70B-Instruct
 project_name: autotrain-llama3-70b-math-v1
 log: tensorboard
-backend: local-cli
+backend: local
 
 data:
   path: rishiraj/guanaco-style-metamath-40k

configs/llm_finetuning/llama3-8b-orpo-space.yml

@@ -0,0 +1,36 @@
+task: llm
+base_model: meta-llama/Meta-Llama-3-8B-Instruct
+project_name: autotrain-llama3-8b-orpo-t1
+log: tensorboard
+backend: spaces-a10g-largex4
+
+data:
+  path: argilla/distilabel-capybara-dpo-7k-binarized
+  train_split: train
+  valid_split: null
+  chat_template: chatml
+  column_mapping:
+    text_column: chosen
+    rejected_text_column: rejected
+
+params:
+  trainer: orpo
+  block_size: 1024
+  model_max_length: 8192
+  max_prompt_length: 512
+  epochs: 3
+  batch_size: 2
+  lr: 3e-5
+  peft: true
+  quantization: int4
+  target_modules: all-linear
+  padding: right
+  optimizer: adamw_torch
+  scheduler: linear
+  gradient_accumulation: 4
+  mixed_precision: fp16
+
+hub:
+  username: ${HF_USERNAME}
+  token: ${HF_TOKEN}
+  push_to_hub: true

configs/llm_finetuning/llama3-8b-orpo.yml

@@ -2,7 +2,7 @@ task: llm
 base_model: meta-llama/Meta-Llama-3-8B-Instruct
 project_name: autotrain-llama3-8b-orpo
 log: tensorboard
-backend: local-cli
+backend: local
 
 data:
   path: argilla/distilabel-capybara-dpo-7k-binarized

configs/text_classification/hub_dataset.yml

@@ -2,7 +2,7 @@ task: text_classification
 base_model: google-bert/bert-base-uncased
 project_name: autotrain-bert-imdb-finetuned
 log: tensorboard
-backend: local-cli
+backend: local
 
 data:
   path: stanfordnlp/imdb

configs/text_classification/local_dataset.yml

@@ -2,7 +2,7 @@ task: text_classification
 base_model: google-bert/bert-base-uncased
 project_name: autotrain-bert-imdb-finetuned
 log: tensorboard
-backend: local-cli
+backend: local
 
 data:
   path: data/ # this must be the path to the directory containing the train and valid files

configs/text_regression/hub_dataset.yml

@@ -2,7 +2,7 @@ task: text_regression
 base_model: google-bert/bert-base-uncased
 project_name: autotrain-bert-sms-spam-finetuned
 log: tensorboard
-backend: local-cli
+backend: local
 
 data:
   path: sms_spam

configs/text_regression/local_dataset.yml

@@ -2,7 +2,7 @@ task: text_regression
 base_model: google-bert/bert-base-uncased
 project_name: autotrain-bert-custom-finetuned
 log: tensorboard
-backend: local-cli
+backend: local
 
 data:
   path: data/ # this must be the path to the directory containing the train and valid files

configs/token_classification/hub_dataset.yml

@@ -2,7 +2,7 @@ task: token_classification
 base_model: google-bert/bert-base-uncased
 project_name: autotrain-bert-conll2003-finetuned
 log: tensorboard
-backend: local-cli
+backend: local
 
 data:
   path: conll2003

configs/token_classification/local_dataset.yml

@@ -2,7 +2,7 @@ task: token_classification
 base_model: google-bert/bert-base-uncased
 project_name: autotrain-bert-custom-finetuned
 log: tensorboard
-backend: local-cli
+backend: local
 
 data:
   path: data/ # this must be the path to the directory containing the train and valid files

src/autotrain/api.py

@@ -1,13 +1,14 @@
 import asyncio
 import os
 import signal
+import sys
 from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
 
 from autotrain import logger
-from autotrain.app_utils import get_running_jobs, run_training
 from autotrain.db import AutoTrainDB
+from autotrain.utils import get_running_jobs, run_training
 
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
@@ -20,6 +21,15 @@
 DB = AutoTrainDB("autotrain.db")
 
 
+def sigint_handler(signum, frame):
+    """Handle SIGINT signal gracefully."""
+    logger.info("SIGINT received. Exiting gracefully...")
+    sys.exit(0)  # Exit with code 0
+
+
+signal.signal(signal.SIGINT, sigint_handler)
+
+
 class BackgroundRunner:
     async def run_main(self):
         while True:

src/autotrain/app.py

@@ -11,7 +11,8 @@
 from huggingface_hub import repo_exists
 from nvitop import Device
 
-from autotrain import __version__, app_utils, logger
+import autotrain.utils as app_utils
+from autotrain import __version__, logger
 from autotrain.app_params import AppParams
 from autotrain.dataset import AutoTrainDataset, AutoTrainDreamboothDataset, AutoTrainImageClassificationDataset
 from autotrain.db import AutoTrainDB
@@ -461,7 +462,7 @@ async def handle_form(
         valid_split = None
 
     logger.info(f"hardware: {hardware}")
-    if hardware == "Local":
+    if hardware == "local-ui":
         running_jobs = app_utils.get_running_jobs(DB)
         if running_jobs:
             raise HTTPException(
@@ -514,7 +515,7 @@ async def handle_form(
                 username=autotrain_user,
                 valid_data=validation_files[0] if validation_files else None,
                 percent_valid=None,  # TODO: add to UI
-                local=hardware.lower() == "local",
+                local=hardware.lower() == "local-ui",
             )
         elif task == "dreambooth":
             dset = AutoTrainDreamboothDataset(
@@ -523,7 +524,7 @@ async def handle_form(
                 token=token,
                 project_name=project_name,
                 username=autotrain_user,
-                local=hardware.lower() == "local",
+                local=hardware.lower() == "local-ui",
             )
 
         else:
@@ -562,7 +563,7 @@ async def handle_form(
                 column_mapping=column_mapping,
                 valid_data=validation_files,
                 percent_valid=None,  # TODO: add to UI
-                local=hardware.lower() == "local",
+                local=hardware.lower() == "local-ui",
                 ext=file_extension,
             )
             if task in ("text-classification", "token-classification"):
@@ -588,16 +589,15 @@ async def handle_form(
     project = AutoTrainProject(params=params, backend=hardware)
     job_id = project.create()
     monitor_url = ""
-    if hardware == "Local":
+    if hardware == "local-ui":
         DB.add_job(job_id)
         monitor_url = "Monitor your job locally / in logs"
-    elif hardware.startswith("EP"):
+    elif hardware.startswith("ep-"):
         monitor_url = f"https://ui.endpoints.huggingface.co/{autotrain_user}/endpoints/{job_id}"
-    else:
+    elif hardware.startswith("spaces-"):
         monitor_url = f"https://hf.co/spaces/{job_id}"
-
-    if job_id is None:
-        monitor_url = "Success! Monitor your job in logs"
+    else:
+        monitor_url = "Success! Monitor your job in logs. Job ID: {job_id}"
 
     return {"success": "true", "monitor_url": monitor_url}