Simple Model Server

This repository provides a FastAPI-based backend for serving machine learning models for classification, entity extraction (NER), and translation (Arabic, Russian, Chinese → English). It is designed for efficient, production-ready deployment with lazy model loading and memory management.

Features

• Text Classification: Zero-shot classification using HuggingFace Transformers.
• Entity Extraction: Named Entity Recognition (NER) with aggregation.
• Translation: MarianMT-based translation for Arabic, Russian, and Chinese to English.
• API Key Authentication: Simple API key header-based security.
• CORS Support: Configurable for cross-origin requests.
• Docker Support: Lightweight Dockerfile for deployment.
• Async Model Loading: Models are loaded on demand and can be unloaded to manage memory.

Requirements

• Python 3.12+
• PyTorch
• Transformers
• FastAPI
• uv (for dependency management)
• Docker (optional, for containerized deployment)

Installation

1. Clone the repository:

   git clone https://github.com/your-org/simple-model-server.git
   cd simple-model-server

2. Set up a virtual environment:

   python3.12 -m venv .venv
   source .venv/bin/activate

   If working with uv, the .venv should be created automatically when you run the next step.

3. Install dependencies:

   pip install -r requirements.txt

   Or, if using uv:

   uv add -r requirements.txt

4. Configure environment variables:

   • Create a .env file (see below for required variables).

Environment Variables

Set the following variables in your .env file or environment:

MODEL_DIRECTORY=/models
API_KEY_1=your_api_key_here
CLASSIFICATION_MODEL_NAME=theme-classification
NER_MODEL_PATH=bert-base-NER
AR_EN_MODEL_NAME=opus-mt-ar-en
RU_EN_MODEL_NAME=opus-mt-ru-en
ZH_EN_MODEL_NAME=opus-mt-zh-en
ROOT_PATH=

5. Storage of models
   • I downloaded all the relevant model files from huggingface and store them in my model directory '/models'. Make sure the names of the models match with the environmenal variables you choose.

Running the Server

Locally

uvicorn app.main:app --host 0.0.0.0 --port 5001 --reload

Or if running with UV

uv run fastapi dev --port 5001

With Docker

Build and run the Docker image:

docker buildx build -t simple-model-server .
docker run -p 5001:5001 --env-file .env simple-model-server

API Endpoints

Health Check

• GET /
  Returns server status, uptime, and memory info.

Classification

• POST /classification/classify
  Classifies text into provided themes.

• POST /classification/classify-multi
  Multi-label classification.

Entity Extraction

• POST /entity/entity
  Extracts named entities from text.

Translation

• POST /translation/translate
  Translates text from Arabic, Russian, or Chinese to English.

Example Request

POST /translation/translate
{
  "text": "مرحبا",
  "source_language": "ARABIC"
}

Authentication

All endpoints require an API key via the X-API-Key header.

X-API-Key: your_api_key_here

Testing

Run tests with:

pytest tests/

Project Structure

app/
  api/                # FastAPI routers and endpoints
  models/             # Model classes (classification, entity, translation)
  middleware/         # Custom middleware (auth, error handling)
  utils/              # Utility modules
  state.py            # Application state and model management
  main.py             # FastAPI app entrypoint
tests/                # Unit and integration tests
Dockerfile            # Docker build instructions
pyproject.toml        # Project dependencies