initial commit

.gitignore

@@ -0,0 +1,7 @@
+# local config
+docker-compose.override.yaml
+
+# PhpStorm / IDEA
+.idea
+# NetBeans
+nbproject

Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.9
+
+RUN useradd -m -u 1000 user
+USER user
+
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH
+
+WORKDIR $HOME/app
+
+COPY --chown=user requirements.txt requirements.txt
+
+RUN pip install --upgrade -r requirements.txt
+
+COPY --chown=user . .
+
+CMD ["uvicorn", "src.main:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -9,3 +9,4 @@ license: other
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+

docker-compose.yaml

@@ -0,0 +1,15 @@
+services:
+  server:
+    build:
+      context: .
+    ports:
+      - 7860:7860
+    develop:
+      watch:
+        - action: rebuild
+          path: .
+    volumes:
+      - python-cache:/home/user/.cache
+
+volumes:
+  python-cache:

requirements.txt

@@ -0,0 +1,13 @@
+fastapi==0.111.*
+requests==2.*
+uvicorn[standard]==0.30.*
+pandas
+transformers 
+datasets 
+evaluate 
+accelerate 
+pillow 
+torchvision 
+scikit-learn
+huggingface_hub
+pydantic

src/abstract_trainer.py

@@ -0,0 +1,20 @@
+from abc import ABC, abstractmethod
+
+import logging
+from .training_status import TrainingStatus
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+
+class AbstractTrainer(ABC): 
+
+    __training_status: TrainingStatus = TrainingStatus(); 
+
+    @abstractmethod
+    async def start_training(self):
+        logger.info('start abstract trainer training')
+        pass
+
+    def get_status(self) -> TrainingStatus:
+        return self.__training_status

src/environment_variable_checker.py

@@ -0,0 +1,27 @@
+import os
+from fastapi import HTTPException, status
+
+class EnvironmentVariableChecker:
+
+    def validate_environment_variables(self): 
+
+        variables = ['AUTHENTICATION_TOKEN', 'HUGGINGFACE_TOKEN', 'HUGGINGFACE_ORGANIZATION']
+
+        for variable in variables:
+            if os.getenv(variable) is None:
+                raise HTTPException(
+                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                    detail=f"Environment variable {variable} not set, please set the {variable} environment variable",
+                )
+
+    def get_authentication_token(self):
+        return os.getenv('AUTHENTICATION_TOKEN')
+    
+    def get_huggingface_token(self): 
+        return os.getenv('HUGGINGFACE_TOKEN'); 
+                         
+    def get_huggingface_organization(self):
+        return os.getenv('HUGGINGFACE_ORGANIZATION');
+
+
+

src/image_classification/image_classification_parameters.py

@@ -0,0 +1,56 @@
+from pydantic import BaseModel
+from typing import Annotated
+from fastapi import Form
+
+
+class ImageClassificationTrainingParameters(BaseModel):
+    epochs: int
+    learning_rate: float
+
+
+def map_image_classification_training_parameters(
+    epocs: Annotated[int, Form(...)] = 3,
+    learning_rate: Annotated[float, Form(...)] = 5e-5
+) -> ImageClassificationTrainingParameters:
+    return ImageClassificationTrainingParameters(
+        epochs=epocs,
+        learning_rate=learning_rate
+    )
+
+
+class ImageClassificationParameters: 
+
+    __training_files_path: str
+    __training_zip_file_path: str
+    __result_model_name: str
+    __source_model_name: str
+    __training_parameters: ImageClassificationTrainingParameters
+
+    def __init__(self, 
+                 training_files_path: str, 
+                 training_zip_file_path: str, 
+                 result_model_name: str, 
+                 source_model_name: str,
+                 training_parameters: ImageClassificationTrainingParameters
+                 ):
+        self.__training_files_path = training_files_path
+        self.__training_zip_file_path = training_zip_file_path
+        self.__result_model_name = result_model_name
+        self.__source_model_name = source_model_name
+        self.__training_parameters = training_parameters
+
+    def get_training_files_path(self) -> str:
+        return self.__training_files_path
+
+    def get_training_zip_file(self) -> str:
+        return self.__training_zip_file_path
+
+    def get_result_model_name(self) -> str:
+        return self.__result_model_name
+
+    def get_source_model_name(self) -> str:
+        return self.__source_model_name
+    
+    def get_training_parameters(self) -> ImageClassificationTrainingParameters:
+        return self.__training_parameters
+    

src/image_classification/image_classification_trainer.py

@@ -0,0 +1,196 @@
+import logging
+
+from ..progress_callback import ProgressCallback
+from ..abstract_trainer import AbstractTrainer
+from ..environment_variable_checker import EnvironmentVariableChecker
+from .image_classification_parameters import ImageClassificationParameters
+
+import zipfile
+import os
+import shutil
+
+from datasets import load_dataset
+from transformers import AutoImageProcessor, DefaultDataCollator, AutoModelForImageClassification, TrainingArguments, Trainer, TrainerState, TrainerControl 
+from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor
+from huggingface_hub import HfFolder
+
+import evaluate
+import numpy as np
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+
+class ImageClassificationTrainer(AbstractTrainer): 
+
+    def start_training(self, parameters: ImageClassificationParameters):
+
+        logger.info('Start Training...')
+
+        try:
+            task = 'Extract training data'
+            self.get_status().update_status(0, task)
+            logger.info(task)
+
+            self.__extract_training_data(parameters)
+
+            if(self.get_status().is_training_aborted()):
+                return
+
+            task = 'Prepare Data set'
+            self.get_status().update_status(10, task)
+            logger.info(task)
+
+            images = self.__prepare_data_set(parameters)
+
+            if(self.get_status().is_training_aborted()):
+                return
+
+            task = 'Start training model'
+            self.get_status().update_status(20, task)
+            logger.info(task)
+
+            self.__train_model(images, parameters)
+
+            self.get_status().update_status(100, "Training completed")
+
+        except Exception as e:
+            logger.error(e)
+            self.get_status().finalize_abort_training(str(e))
+
+            raise RuntimeError(f"An error occurred: {str(e)}")
+
+        finally:
+            # Cleanup after processing
+            logger.info('Cleaning up training files after training')
+            shutil.rmtree(parameters.get_training_files_path())    
+
+            if(self.get_status().is_training_aborted()):
+                self.get_status().finalize_abort_training("Training aborted")
+                
+
+
+    def __extract_training_data(self, parameters: ImageClassificationParameters):
+        training_file = parameters.get_training_zip_file()
+
+        # Check if it is a valid ZIP file
+        if not zipfile.is_zipfile(training_file):
+            raise RuntimeError("Uploaded file is not a valid zip file")
+
+        # Extract the ZIP file
+        with zipfile.ZipFile(training_file, 'r') as zip_ref:
+            zip_ref.extractall(parameters.get_training_files_path())
+
+        os.remove(training_file)
+        logger.info(os.listdir(parameters.get_training_files_path()))
+
+
+
+    def __prepare_data_set(self, parameters: ImageClassificationParameters) -> dict: 
+
+        dataset = load_dataset("imagefolder", data_dir=parameters.get_training_files_path())
+
+        images = dataset["train"]
+        images = images.train_test_split(test_size=0.2)
+        
+        logger.info(images)
+        logger.info(images["train"][100])
+
+
+        # Preprocess the images
+        image_processor = AutoImageProcessor.from_pretrained(parameters.get_source_model_name())
+
+        # Apply some image transformations to the images to make the model more robust against overfitting.
+        normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
+        size = (
+            image_processor.size["shortest_edge"]
+            if "shortest_edge" in image_processor.size
+            else (image_processor.size["height"], image_processor.size["width"])
+        )
+        _transforms = Compose([RandomResizedCrop(size), ToTensor(), normalize])
+
+        def transforms(examples):
+            examples["pixel_values"] = [_transforms(img.convert("RGB")) for img in examples["image"]]
+            del examples["image"]
+            return examples
+
+        images = images.with_transform(transforms)
+
+        return images
+
+
+    def __train_model(self, images: dict, parameters: ImageClassificationParameters):
+
+        environment_variable_checker = EnvironmentVariableChecker()
+        HfFolder.save_token(environment_variable_checker.get_huggingface_token())
+        
+        image_processor = AutoImageProcessor.from_pretrained(parameters.get_source_model_name())
+        data_collator = DefaultDataCollator()
+        progressCallback = ProgressCallback(self.get_status())
+
+        # Evaluate and metrics
+        accuracy = evaluate.load("accuracy")
+        def compute_metrics(eval_pred):
+            predictions, labels = eval_pred
+            predictions = np.argmax(predictions, axis=1)
+            return accuracy.compute(predictions=predictions, references=labels)
+
+        # get label maps
+        labels = images["train"].features["label"].names
+        label2id, id2label = dict(), dict()
+        for i, label in enumerate(labels):
+            label2id[label] = str(i)
+            id2label[str(i)] = label        
+        logger.info(id2label)
+
+        # train the model
+        model = AutoModelForImageClassification.from_pretrained(
+            parameters.get_source_model_name(),
+            num_labels=len(labels),
+            id2label=id2label,
+            label2id=label2id,
+        )
+
+        target_model_id = environment_variable_checker.get_huggingface_organization() + '/' + parameters.get_result_model_name()
+        training_args = TrainingArguments(
+            output_dir=parameters.get_result_model_name(),
+            hub_model_id=target_model_id,
+            remove_unused_columns=False,
+            eval_strategy="epoch",
+            save_strategy="epoch",
+            learning_rate=parameters.get_training_parameters().learning_rate,
+            per_device_train_batch_size=16,
+            gradient_accumulation_steps=4,
+            per_device_eval_batch_size=16,
+            num_train_epochs=parameters.get_training_parameters().epochs,
+            warmup_ratio=0.1,
+            logging_steps=10,
+            load_best_model_at_end=True,
+            metric_for_best_model="accuracy",
+            push_to_hub=False,
+            hub_private_repo=True,
+        )
+
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            data_collator=data_collator,
+            train_dataset=images["train"],
+            eval_dataset=images["test"],
+            tokenizer=image_processor,
+            compute_metrics=compute_metrics,
+            callbacks=[progressCallback]
+        )
+
+
+        if(self.get_status().is_training_aborted()):
+            return
+
+        trainer.train()
+
+        if(self.get_status().is_training_aborted()):
+            return
+
+        logger.info(f"Model trained, start uploading")
+        self.get_status().update_status(90, f"Uploading model to Hugging Face")
+        trainer.push_to_hub()

src/main.py

@@ -0,0 +1,145 @@
+import os
+import requests
+import torch
+
+from .training_status import Status
+from .environment_variable_checker import EnvironmentVariableChecker
+
+from .task_manager import TaskManager
+from .training_manager import TrainingManager
+from .image_classification.image_classification_trainer import ImageClassificationTrainer
+from .image_classification.image_classification_parameters import ImageClassificationParameters, map_image_classification_training_parameters, ImageClassificationTrainingParameters 
+
+from fastapi import FastAPI, Header, Depends, HTTPException, BackgroundTasks, UploadFile, Form, File, status
+from fastapi.responses import FileResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from pydantic import BaseModel
+from typing import Optional, Annotated
+
+
+import logging
+import sys
+
+import zipfile
+import os
+from pathlib import Path
+import tempfile
+import shutil
+
+
+app = FastAPI()
+
+environmentVariableChecker = EnvironmentVariableChecker()
+environmentVariableChecker.validate_environment_variables()
+
+logging.basicConfig(format='%(asctime)s %(levelname)-8s %(message)s')
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+classification_trainer: TrainingManager = TrainingManager(ImageClassificationTrainer())
+
+security = HTTPBearer()
+def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
+
+    token = environmentVariableChecker.get_authentication_token()
+    
+    if credentials.credentials != token:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid token",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    return {"token": credentials.credentials}
+
+
+class ResponseModel(BaseModel):
+    message: str
+    success: bool = True
+
+
+@app.post(
+    "/upload", 
+    summary="Upload a zip file containing training data",
+    response_model=ResponseModel
+)
+async def upload_file(
+    training_params: Annotated[ImageClassificationTrainingParameters, Depends(map_image_classification_training_parameters)],
+    data_files_training: Annotated[UploadFile, File(...)],
+    token_data: dict = Depends(verify_token),
+    result_model_name: str = Form(...),
+    source_model_name: str = Form('google/vit-base-patch16-224-in21k'),
+):
+
+    # check if training is running, if so then exit
+    status = classification_trainer.get_task_status()
+    if status.get_status() == Status.IN_PROGRESS or status.get_status() == Status.CANCELLING:
+        raise HTTPException(status_code=405, detail="Training is already in progress")
+
+    # Ensure the uploaded file is a ZIP file
+    if not data_files_training.filename.endswith(".zip"):
+        raise HTTPException(status_code=422, detail="Uploaded file is not a zip file")
+
+    try:
+        # Create a temporary directory to extract the contents
+        tmp_path = os.path.join(tempfile.gettempdir(), 'training_data')
+        path = Path(tmp_path)
+        path.mkdir(parents=True, exist_ok=True)
+
+        contents = await data_files_training.read()
+        zip_path = os.path.join(tmp_path, 'image_classification_data.zip')
+        with open(zip_path, 'wb') as temp_file:
+            temp_file.write(contents)
+
+        # prepare parameters
+        parameters = ImageClassificationParameters(
+            training_files_path=tmp_path,
+            training_zip_file_path=zip_path,
+            result_model_name=result_model_name,
+            source_model_name=source_model_name,
+            training_parameters=training_params
+        )
+
+        # start training
+        await classification_trainer.start_training(parameters)
+        
+        # TODO add more return parameters and information
+        return ResponseModel(message="training started")
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+
+
+@app.get("/get_task_status")
+async def get_task_status(token_data: dict = Depends(verify_token)):
+    status = classification_trainer.get_task_status()
+    return {
+        "progress": status.get_progress(),
+        "task": status.get_task(),
+        "status": status.get_status().value
+    }
+
+
+@app.get("/stop_task")
+async def stop_task(token_data: dict = Depends(verify_token)):
+    try: 
+        classification_trainer.stop_task()
+        return {
+            "success": True
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+
+
+
+@app.get("/gpu_check")
+async def gpu_check():
+
+    gpu = 'GPU not available'
+    if torch.cuda.is_available():
+        gpu = 'GPU is available'
+        print("GPU is available")
+    else:
+        print("GPU is not available")
+
+    return {'success': True, 'response': 'hello world 3', 'gpu': gpu}
+

src/progress_callback.py

@@ -0,0 +1,31 @@
+import logging
+from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl
+
+from .training_status import TrainingStatus
+
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+class ProgressCallback(TrainerCallback):
+
+    __trainingStatus: TrainingStatus = None
+
+    def __init__(self, trainingStatus: TrainingStatus):
+        self.__trainingStatus = trainingStatus
+
+    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        logger.info(f"Completed step {state.global_step} of {state.max_steps}")
+
+        if self.__trainingStatus.is_training_aborted(): 
+            control.should_training_stop = True
+            logger.info("Training aborted")
+            return
+
+        startPercentage = 21
+        endPercentage = 89
+        scope = endPercentage - startPercentage
+        progress = startPercentage + (state.global_step / state.max_steps) * scope
+
+        self.__trainingStatus.update_status(progress, f"Training model, completed step {state.global_step} of {state.max_steps}")
+

src/task_manager.py

@@ -0,0 +1,72 @@
+import asyncio
+import logging
+from fastapi import BackgroundTasks, HTTPException
+
+from concurrent.futures import ThreadPoolExecutor
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+
+class Worker: 
+    def doing_work(self, task_manager):
+        task_manager.task_status["status"] = "Running"
+        for i in range(1, 101):
+            if task_manager.task_status["status"] == "Stopped":
+                break
+            asyncio.sleep(1)  # Simulate a time-consuming task
+            task_manager.task_status["progress"] = i
+            logger.info('process ' + str(i) + '%' + ' done')
+
+        if task_manager.task_status["status"] != "Stopped":
+            task_manager.task_status["status"] = "Completed"
+
+
+
+class TaskManager:
+
+    task_status = {"progress": 0, "status": "Not started"}
+    task = None
+
+    #def __init__(self):
+
+    worker = Worker()
+
+    async def doing_work(self):
+        loop = asyncio.get_running_loop()
+        with ThreadPoolExecutor() as pool:
+            await loop.run_in_executor(pool, self.worker.doing_work, self)
+            #self.worker.doing_work(self)    
+
+        # self.task_status["status"] = "Running"
+        # for i in range(1, 101):
+        #     if self.task_status["status"] == "Stopped":
+        #         break
+        #     await asyncio.sleep(1)  # Simulate a time-consuming task
+        #     self.task_status["progress"] = i
+        #     logger.info('process ' + str(i) + '%' + ' done')
+
+        # if self.task_status["status"] != "Stopped":
+        #     self.task_status["status"] = "Completed"
+
+
+    async def start_task(self):
+        if self.task is None or self.task.done():
+            self.task_status["progress"] = 0
+            self.task_status["status"] = "Not started"
+            self.task = asyncio.create_task(self.doing_work())
+            return {"message": "Task started"}
+        else:
+            raise HTTPException(status_code=409, detail="Task already running")
+
+    async def get_task_status(self):
+        return self.task_status
+
+    async def stop_task(self):
+        if self.task is not None and not self.task.done():
+            self.task_status["status"] = "Stopped"
+            self.task.cancel()
+            return {"message": "Task stopped"}
+        else:
+            raise HTTPException(status_code=409, detail="No task running")
+

src/training_manager.py

@@ -0,0 +1,50 @@
+import logging
+import asyncio
+from .abstract_trainer import AbstractTrainer
+from .training_status import TrainingStatus
+from concurrent.futures import ThreadPoolExecutor
+
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+
+
+class TrainingManager: 
+
+    __training_task = None
+    __trainer: AbstractTrainer = None
+
+    task_status = {"progress": 0, "status": "Not started"}
+
+    def __init__(self, trainer: AbstractTrainer):
+        self.__trainer = trainer
+
+    async def __do_start_training(self, parameters):
+        logger.info('do start training')
+        
+        loop = asyncio.get_running_loop()
+        with ThreadPoolExecutor() as pool:
+            await loop.run_in_executor(pool, self.__trainer.start_training, parameters)
+
+        logger.info('done')
+
+    async def start_training(self, parameters):
+        logger.info('start training')
+
+        if self.__training_task is None or self.__training_task.done():
+            self.__training_task = asyncio.create_task(self.__do_start_training(parameters))
+        else:
+            raise RuntimeError("Training already running")
+
+    def get_task_status(self) -> TrainingStatus:
+        return self.__trainer.get_status()
+
+    def stop_task(self):
+        if self.__training_task is not None and not self.__training_task.done():
+            self.__trainer.get_status().abort_training("Stopping training")
+            #self.__training_task.cancel()
+
+        else:
+            raise RuntimeError("No task running")
+

src/training_status.py

@@ -0,0 +1,58 @@
+import logging
+from enum import Enum
+
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+class Status(Enum):
+    NOT_STARTED = "NOT_STARTED"
+    IN_PROGRESS = "IN_PROGRESS"
+    CANCELLING = "CANCELLING"
+    CANCELLED = "CANCELLED"
+    COMPLETED = "COMPLETED"
+
+class TrainingStatus: 
+
+    __status: Status = Status.NOT_STARTED
+    __task: str = None
+    __progress: int = 0
+
+    def update_status(self, progress: int, task: str):
+        if progress < 0 or progress > 100:
+            raise ValueError("Progress must be between 0 and 100")
+        
+        if progress > 0:
+            self.__status = Status.IN_PROGRESS    
+
+        if progress == 100:
+            self.__status = Status.COMPLETED
+
+        self.__progress = progress
+        
+        if task is not None:
+            self.__task = task
+
+    
+    def abort_training(self, task: str):
+        self.__task = task
+        self.__status = Status.CANCELLING
+
+    def finalize_abort_training(self, task: str):
+        self.__status = Status.CANCELLED
+        self.__progress = 0
+        self.__task = task
+
+    def is_training_aborted(self) -> bool:
+        return (self.__status == Status.CANCELLING)
+    
+    def get_status(self) -> str:
+        return self.__status
+    
+    def get_progress(self) -> int:
+        return self.__progress
+    
+    def get_task(self) -> str:
+        return self.__task
+
+