additional tasks

requirements.txt

@@ -5,4 +5,5 @@ transformers
 sentencepiece
 sacremoses
 torch
+pillow
 # Optional dependencies for specific features

src/main.py

@@ -11,26 +11,15 @@
 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 import FastAPI, Path, Depends, HTTPException, UploadFile, Form, File, status, Request
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 from typing import Annotated
-
+import json
 
 import logging
-from pathlib import Path
-import tempfile
 import sys
+import base64
 
 
 from transformers import pipeline
@@ -41,9 +30,6 @@ app = FastAPI(
     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)
@@ -65,7 +51,6 @@ class StreamToLogger(object):
 sys.stdout = StreamToLogger(logger, logging.INFO)
 sys.stderr = StreamToLogger(logger, logging.ERROR)
 
-#classification_trainer: TrainingManager = TrainingManager()
 
 
 class ResponseModel(BaseModel):
@@ -74,51 +59,6 @@ class ResponseModel(BaseModel):
     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 """
@@ -133,39 +73,73 @@ async def gpu_check():
     return {'success': True, 'gpu': gpu}
 
 
-from fastapi import Body
 from typing import Optional
 
+
+
+# =========================
+# Translation Task
+# =========================
+
 class TranslationRequest(BaseModel):
     inputs: str
     parameters: Optional[dict] = None
+    options: Optional[dict] = None
+
+async def get_translation_request(
+    request: Request
+)  -> TranslationRequest: 
+    content_type = request.headers.get("content-type", "")
+    if content_type.startswith("application/json"):
+        data = await request.json()
+        return TranslationRequest(**data)
+    if content_type.startswith("application/x-www-form-urlencoded"):
+        raw = await request.body()
+        try:
+            data = json.loads(raw)
+            return TranslationRequest(**data)
+        except Exception:
+            try:
+                data = json.loads(raw.decode("utf-8"))
+                return TranslationRequest(**data)
+            except Exception:
+                raise HTTPException(status_code=400, detail="Invalid request body")
+    raise HTTPException(status_code=400, detail="Unsupported content type")
+
+
 
 @app.post(
     "/translation/{model_name:path}/", 
-)
-async def translation(
-    model_name: str,
-    body: TranslationRequest = Body(
-        ...,
-        example={
-            "inputs": "I am a car",
-            "parameters": {
-                "repetition_penalty": 1.6,
+    openapi_extra={
+        "requestBody": {
+            "content": {
+                "application/json": {
+                    "example": {
+                        "inputs": "Hello, world! foo bar",
+                        "parameters": {"repetition_penalty": 1.6}
+                    }
+                }
             }
         }
+    }        
+)
+async def translate(
+    request: Request,
+    model_name: str = Path(
+        ...,
+        description="The name of the translation model (e.g. Helsinki-NLP/opus-mt-en-de)",
+        example="Helsinki-NLP/opus-mt-en-de"
     )
-):
+    ):
     """
     Execute translation tasks.
 
-    Args:
-        model_name (str): The HuggingFace model name to use for translation.
-        body (TranslationRequest): The request payload containing translation parameters.
-
     Returns:
         list: The translation result(s) as returned by the pipeline.
     """
 
+    translationRequest: TranslationRequest = await get_translation_request(request)
+
     try:
         pipe = pipeline("translation", model=model_name)
     except Exception as e:
@@ -176,7 +150,96 @@ async def translation(
         )
 
     try:       
-        result = pipe(body.inputs, **(body.parameters or {}))
+        result = pipe(translationRequest.inputs, **(translationRequest.parameters or {}))
+    except Exception as e:
+        logger.error(f"Inference failed for model '{model_name}': {str(e)}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Inference failed: {str(e)}"
+        )
+
+    return result
+
+
+# =========================
+# Zero-Shot Image Classification Task
+# =========================
+
+
+class ZeroShotImageClassificationRequest(BaseModel):
+    inputs: str
+    parameters: Optional[dict] = None
+
+async def get_zero_shot_image_classification_request(
+    request: Request
+)  -> ZeroShotImageClassificationRequest: 
+    content_type = request.headers.get("content-type", "")
+    if content_type.startswith("application/json"):
+        data = await request.json()
+        return ZeroShotImageClassificationRequest(**data)
+    if content_type.startswith("application/x-www-form-urlencoded"):
+        raw = await request.body()
+        try:
+            data = json.loads(raw)
+            return ZeroShotImageClassificationRequest(**data)
+        except Exception:
+            try:
+                data = json.loads(raw.decode("utf-8"))
+                return ZeroShotImageClassificationRequest(**data)
+            except Exception:
+                raise HTTPException(status_code=400, detail="Invalid request body")
+    raise HTTPException(status_code=400, detail="Unsupported content type")
+
+
+
+@app.post(
+    "/zero-shot-image-classification/{model_name:path}/",
+    openapi_extra={
+        "requestBody": {
+            "content": {
+                "application/json": {
+                    "example": {
+                        "inputs": "base64_encoded_image_string",
+                        "parameters": {"candidate_labels": "green, yellow, blue, white, silver"}
+                    }
+                }
+            }
+        }        
+    }
+)
+async def zero_shot_image_classification(
+    request: Request,
+    model_name: str = Path(
+        ...,
+        description="The name of the zero-shot classification model (e.g., openai/clip-vit-large-patch14-336)",
+        example="openai/clip-vit-large-patch14-336"
+    )
+    ):
+    """
+    Execute zero-shot image classification tasks.
+
+    Returns:
+        list: The classification result(s) as returned by the pipeline.
+    """
+
+    zeroShotRequest: ZeroShotImageClassificationRequest = await get_zero_shot_image_classification_request(request)
+
+    try:
+        pipe = pipeline("zero-shot-image-classification", model=model_name)
+    except Exception as e:
+        logger.error(f"Failed to load model '{model_name}': {str(e)}")
+        raise HTTPException(
+            status_code=404,
+            detail=f"Model '{model_name}' could not be loaded: {str(e)}"
+        )
+
+    try:       
+        candidate_labels = []
+        if zeroShotRequest.parameters:
+            candidate_labels = zeroShotRequest.parameters.get('candidate_labels', [])
+            if isinstance(candidate_labels, str):
+                candidate_labels = [label.strip() for label in candidate_labels.split(',')]
+        result = pipe(zeroShotRequest.inputs, candidate_labels=candidate_labels)
     except Exception as e:
         logger.error(f"Inference failed for model '{model_name}': {str(e)}")
         raise HTTPException(
@@ -187,117 +250,85 @@ async def translation(
     return result
 
 
-# ===========================================
-# 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)}")
 
+# =========================
+# Image to Text Task
+# =========================
+
+
+async def get_encoded_image(
+    request: Request
+)  -> str: 
+    content_type = request.headers.get("content-type", "")
+    if content_type.startswith("multipart/form-data"):
+        form = await request.form()
+        image = form.get("image")
+        if image:
+            image_bytes = await image.read()
+            return base64.b64encode(image_bytes).decode("utf-8")
+    if content_type.startswith("image/"):
+        image_bytes = await request.body()
+        return base64.b64encode(image_bytes).decode("utf-8")
+
+    raise HTTPException(status_code=400, detail="Unsupported content type")
+
+
+
+@app.post(
+    "/image-to-text/{model_name:path}/",
+    openapi_extra={
+        "requestBody": {
+            "content": {
+                "multipart/form-data": {
+                    "schema": {
+                        "type": "object",
+                        "properties": {
+                            "image": {
+                                "type": "string",
+                                "format": "binary",
+                                "description": "Image file to upload"
+                            }
+                        },
+                        "required": ["image"]
+                    }
+                }
+            }
+        }
+    }
+)
+async def image_to_text(
+    request: Request,
+    model_name: str = Path(
+        ...,
+        description="The name of the image-to-text (e.g., Salesforce/blip-image-captioning-base)",
+        example="Salesforce/blip-image-captioning-base"
+    )
+    ):
+    """
+    Execute image-to-text tasks.
+
+    Returns:
+        list: The generated text as returned by the pipeline.
+    """
+
+    encoded_image = await get_encoded_image(request)
+
+    try:
+        pipe = pipeline("image-to-text", model=model_name, use_fast=True)
+    except Exception as e:
+        logger.error(f"Failed to load model '{model_name}': {str(e)}")
+        raise HTTPException(
+            status_code=404,
+            detail=f"Model '{model_name}' could not be loaded: {str(e)}"
+        )
+
+    try:       
+        result = pipe(encoded_image)
+    except Exception as e:
+        logger.error(f"Inference failed for model '{model_name}': {str(e)}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Inference failed: {str(e)}"
+        )
+
+    return result