Offline Rest Translation Service

python main.py

curl 'http://localhost:5080/translate' \
  -X POST \
  -H 'Accept: application/json' \
  -H 'Content-Type: application/json' \
  --data-raw '{"source":"en","target":"sw","input":["Hello, World!","How are you?"]}'

{
  "translation": [
    "Halo, Ulimwengu!",
    "Unaendeleaje?"
  ]
}

import os
import logging

from functools import lru_cache
from flask import Flask, request, jsonify, abort
from transformers import MarianMTModel, MarianTokenizer
from flask_cors import CORS

# Configure logging
logging.basicConfig(level=logging.INFO)

# Initialize Flask app
app = Flask('demsking-translate')
CORS(app)  # enable CORS for all routes

@lru_cache(maxsize=128)  # cache up to 128 entries
def get_model(src, tgt):
  # Load the pre-trained model and tokenizer for the requested language pair
  model_name = f"Helsinki-NLP/opus-mt-{src}-{tgt}"
  
  logging.info(f'Loading model {model_name}')

  model = MarianMTModel.from_pretrained(model_name)
  tokenizer = MarianTokenizer.from_pretrained(model_name)

  return model, tokenizer

# Endpoint for translation service
@app.route('/translate', methods=['POST'])
def translate():
  # Get input text and language codes from request body
  input_text = request.json.get('input')
  source_lang = request.json.get('source')
  target_lang = request.json.get('target')

  # Check if all required parameters are included in request body
  if not all([input_text, source_lang, target_lang]):
    return abort(400, 'Missing required parameters')

  logging.info(f'Translating from "{source_lang}" to "{target_lang}": "{input_text}"')

  try:
    # Load the pre-trained model and tokenizer for the requested language pair
    model, tokenizer = get_model(source_lang, target_lang)

    # Tokenize the input text and convert language codes to model-specific format
    input_ids = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512, add_special_tokens=True).input_ids

    # Generate the output text
    output_ids = model.generate(input_ids)
    output_text = tokenizer.batch_decode(output_ids, skip_special_tokens=True)

    logging.info(f'Translation: {output_text}')

    # Return the output text as JSON response
    return jsonify({'translation': output_text})
  except Exception:
    return abort(400, f'Unable to translate from "{source_lang}" to "{target_lang}"')

# Run the Flask app
if __name__ == '__main__':
  # Start the Flask app
  app.run(host='0.0.0.0', port=5080, debug=True)

{ pkgs ? import <nixpkgs> {} }:

let
  envDir = "$(pwd)/venv";
in

pkgs.mkShell {
  nativeBuildInputs = [
    pkgs.gnumake
    pkgs.python310
    pkgs.python310Packages.flask
    pkgs.python310Packages.flask-cors
    pkgs.python310Packages.torch
    pkgs.python310Packages.transformers
    pkgs.python310Packages.sentencepiece
    pkgs.python310Packages.sacremoses
    pkgs.python310Packages.gunicorn
  ];
  shellHook = ''
    export LD_LIBRARY_PATH=${envDir}/lib
    export VIRTUAL_ENV_DISABLE_PROMPT=true

    virtualenv `basename ${envDir}`

    export PIP_PREFIX=${envDir}
    export PYTHONUSERBASE=${envDir}
    export PYTHON_SITE_PACKAGES=$PIP_PREFIX/${pkgs.python310.sitePackages}
    export PYTHONPATH="$PYTHON_SITE_PACKAGES:$PYTHONPATH"
    export PATH="$PIP_PREFIX/bin:${pkgs.ruff}/bin:$PATH"
    unset SOURCE_DATE_EPOCH

    source ${envDir}/bin/activate
  '';
}