src/main.py

# -------------------------------------------------------------------
# Pimcore
#
# This source file is available under two different licenses:
#  - GNU General Public License version 3 (GPLv3)
#  - Pimcore Commercial License (PCL)
# Full copyright and license information is available in
# LICENSE.md which is distributed with this source code.
#
#  @copyright  Copyright (c) Pimcore GmbH (http://www.pimcore.org)
#  @license    http://www.pimcore.org/license     GPLv3 and PCL
# -------------------------------------------------------------------


import os
import torch

from .training_status import Status
from .environment_variable_checker import EnvironmentVariableChecker

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 .text_classification.text_classification_trainer import TextClassificationTrainer
from .text_classification.text_classification_parameters import TextClassificationParameters, map_text_classification_training_parameters, TextClassificationTrainingParameters 


from fastapi import FastAPI, Depends, HTTPException, UploadFile, Form, File, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel
from typing import Annotated


import logging
import os
from pathlib import Path
import tempfile


app = FastAPI(
    title="Pimcore Fine-Tuning Service",
    description="This service allows you to fine-tune image and text classification models and upload them to hugging face hub.",
    version="1.0.0"
)

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()


class ResponseModel(BaseModel):
    """ Default pesponse model for endpoints. """
    message: str
    success: bool = True


# ===========================================
# Security Check
# ===========================================

security = HTTPBearer()
def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
    """Verify the token provided by the user."""

    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}


# ===========================================
# Training Status Endpoints
# ===========================================

@app.get("/get_training_status")
async def get_task_status(token_data: dict = Depends(verify_token)):
    """ Get the status of the currently running training (if any). """
    status = classification_trainer.get_task_status()
    return {
        "project": status.get_project_name(),
        "progress": status.get_progress(),
        "task": status.get_task(),
        "status": status.get_status().value
    }

@app.put("/stop_training")
async def stop_task(token_data: dict = Depends(verify_token)):
    """ Stop the currently running training (if any). """
    try: 
        status = classification_trainer.get_task_status()
        classification_trainer.stop_task()
        return ResponseModel(message=f"Training stopped for `{ status.get_project_name() }`")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")


@app.get("/gpu_check")
async def gpu_check():
    """ Check if a GPU is available """

    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, 'gpu': gpu}


# ===========================================
# Fine-Tuning Image Classification 
# ===========================================

@app.post(
    "/training/image_classification", 
    response_model=ResponseModel
)
async def image_classification(
    training_params: Annotated[ImageClassificationTrainingParameters, Depends(map_image_classification_training_parameters)],
    training_data_zip: Annotated[UploadFile, File(description="The ZIP file containing the training data, with a folder per class which contains images belonging to that class.")],
    token_data: dict = Depends(verify_token),
    project_name: str = Form(description="The name of the project. Will also be used as name of resulting model that will be created after fine tuning and as the name of the repository at huggingface."),
    source_model_name: str = Form('google/vit-base-patch16-224-in21k', description="The source model to be used as basis for fine tuning."),
):
    """
    Start fine tuning an image classification model with the provided data.
    """

    # 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 training_data_zip.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 training_data_zip.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,
            project_name=project_name,
            source_model_name=source_model_name,
            training_parameters=training_params
        )

        # start training
        await classification_trainer.start_training(ImageClassificationTrainer(), parameters)
        
        return ResponseModel(message="Training started.")
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")




# ===========================================
# Fine-Tuning Text Classification 
# ===========================================

@app.post(
    "/training/text_classification", 
    response_model=ResponseModel
)
async def text_classificaiton(
    training_params: Annotated[TextClassificationTrainingParameters, Depends(map_text_classification_training_parameters)],
    training_data_csv: Annotated[UploadFile, File(description="The CSV file containing the training data, necessary columns `value` (text data) and `target` (classification).")],
    token_data: dict = Depends(verify_token),
    project_name: str = Form(description="The name of the project. Will also be used as name of resulting model that will be created after fine tuning and as the name of the repository at huggingface."),
    training_csv_limiter: str = Form(';', description="The delimiter used in the CSV file."),   
    source_model_name: str = Form('distilbert/distilbert-base-uncased'),
):
    """Start fine tuning an text classification model with the provided data."""

    # 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 CSV file
    if not training_data_csv.filename.endswith(".csv"):
        raise HTTPException(status_code=422, detail="Uploaded file is not a csv 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 training_data_csv.read()
        csv_path = os.path.join(tmp_path, 'data.csv')
        with open(csv_path, 'wb') as temp_file:
            temp_file.write(contents)

        # prepare parameters
        parameters = TextClassificationParameters(
            training_csv_file_path=csv_path,
            training_csv_limiter=training_csv_limiter,
            project_name=project_name,
            source_model_name=source_model_name,
            training_parameters=training_params
        )

        # start training
        await classification_trainer.start_training(TextClassificationTrainer(), parameters)
        
        return ResponseModel(message="Training started.")
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")