Added Translation Endpoint

Dockerfile

@@ -1,6 +1,36 @@
 FROM pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime
 ENV DEBIAN_FRONTEND=noninteractive
 
+RUN apt-get update && \
+    apt-get upgrade -y && \
+    apt-get install -y --no-install-recommends \
+    git \
+    git-lfs \
+    wget \
+    curl \
+    # python build dependencies \
+    build-essential \
+    libssl-dev \
+    zlib1g-dev \
+    libbz2-dev \
+    libreadline-dev \
+    libsqlite3-dev \
+    libncursesw5-dev \
+    xz-utils \
+    tk-dev \
+    libxml2-dev \
+    libxmlsec1-dev \
+    libffi-dev \
+    liblzma-dev \
+    # gradio dependencies \
+    ffmpeg
+
+# fairseq2 dependencies
+RUN apt-get install -y --no-install-recommends \
+    libsndfile-dev
+
+RUN apt-get clean && rm -rf /var/lib/apt/lists/*
+
 RUN useradd -m -u 1000 user
 USER user
 ENV HOME=/home/user \
@@ -15,5 +45,10 @@ RUN pip install -r ${HOME}/app/requirements.txt
 # RUN mkdir content
 # ADD --chown=user https://<SOME_ASSET_URL> content/<SOME_ASSET_NAME>
 
+# SeamlessCommunication requirements
+RUN pip install -r ${HOME}/app/seamless_requirements.txt && \
+    pip install fairseq2 --pre --extra-index-url https://fair.pkg.atmeta.com/fairseq2/pt2.1.0/cu121 && \
+    pip install ${HOME}/app/whl/seamless_communication-1.0.0-py3-none-any.whl
+
 # Start the FastAPI app on port 7860, the default port expected by Spaces
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -22,4 +22,6 @@ Users should be able to call the task and get back in the standard format
     "model": "BAAI/bge-base-en-v1.5",
     "inputs: ["This is one text", "This is second text"],
     "parameters": {}
-}
+}
+
+TODO: Models are cached in volume directory

app.py

@@ -1,6 +1,6 @@
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi import FastAPI
-import sentence_embeddings
+from tasks import sentence_embeddings
 
 app = FastAPI(docs_url="/", redoc_url=None)
 

config.py

@@ -1,6 +1,19 @@
+import torch
 import os
 import dotenv
 
 dotenv.load_dotenv()
 
-TEST_MODE = (os.getenv('TEST_MODE', 'False') == "True")
+TEST_MODE = (os.getenv('TEST_MODE', 'False') == "True")
+
+if torch.cuda.is_available():
+    device = torch.device("cuda:0")
+    dtype = torch.float16
+else:
+    device = torch.device("cpu")
+    dtype = torch.float32
+
+from datetime import datetime
+
+def log(data: dict):
+    print(f"{datetime.now().isoformat()}: {data}")

logger.py

@@ -1,4 +0,0 @@
-from datetime import datetime
-
-def log(data: dict):
-    print(f"{datetime.now().isoformat()}: {data}")

requirements.txt

@@ -2,4 +2,5 @@ transformers
 torch
 fastapi
 uvicorn
+pydantic
 python-dotenv

seamless_requirements.txt

@@ -0,0 +1,2 @@
+omegaconf==2.3.0
+fasttext==0.9.2

tasks/sentence_embeddings.py

@@ -0,0 +1,83 @@
+from typing import Optional
+from fastapi import APIRouter
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModel
+import torch
+from datetime import datetime
+from config import TEST_MODE, device, log
+
+router = APIRouter()
+
+class SentenceEmbeddingsInput(BaseModel):
+    inputs: list[str]
+    model: str
+    parameters: dict
+
+class SentenceEmbeddingsOutput(BaseModel):
+    embeddings: Optional[list[list[float]]] = None
+    error: Optional[str] = None
+
+@router.post('/sentence-embeddings')
+def sentence_embeddings(inputs: SentenceEmbeddingsInput):
+    start_time = datetime.now()
+    fn = sentence_embeddings_mapping.get(inputs.model)
+    if not fn:
+        return SentenceEmbeddingsOutput(
+            error=f'No sentence embeddings model found for {inputs.model}'
+        )
+    
+    try:
+        embeddings = fn(inputs.inputs, inputs.parameters)
+        
+        log({
+            "task": "sentence_embeddings",
+            "model": inputs.model,
+            "start_time": start_time.isoformat(),
+            "time_taken": (datetime.now() - start_time).total_seconds(),
+            "inputs": inputs.inputs,
+            "outputs": embeddings,
+            "parameters": inputs.parameters,
+        })
+        loaded_models_last_updated[inputs.model] = datetime.now()
+        return SentenceEmbeddingsOutput(
+            embeddings=embeddings
+        )
+    except Exception as e:
+        return SentenceEmbeddingsOutput(
+            error=str(e)
+        )
+
+def generic_sentence_embeddings(model_name: str):
+    global loaded_models
+
+    def process_texts(texts: list[str], parameters: dict):
+        if TEST_MODE:
+            return [[0.1,0.2]] * len(texts)
+    
+        if model_name in loaded_models:
+            tokenizer, model = loaded_models[model_name]
+        else:
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            model = AutoModel.from_pretrained(model_name).to(device)
+            loaded_models[model] = (tokenizer, model)
+
+        # Tokenize sentences
+        encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt').to(device)
+        with torch.no_grad():
+            model_output = model(**encoded_input)
+            sentence_embeddings = model_output[0][:, 0]
+
+        # normalize embeddings
+        sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
+        return sentence_embeddings.tolist()
+
+    return process_texts
+
+# Polling every X minutes to 
+loaded_models = {}
+loaded_models_last_updated = {}
+
+sentence_embeddings_mapping = {
+    'BAAI/bge-base-en-v1.5': generic_sentence_embeddings('BAAI/bge-base-en-v1.5'),
+    'BAAI/bge-large-en-v1.5': generic_sentence_embeddings('BAAI/bge-large-en-v1.5'),
+}

tasks/translation.py

@@ -0,0 +1,135 @@
+from fastapi import APIRouter
+from pydantic import BaseModel
+from typing import Optional
+from config import TEST_MODE, device, dtype, log
+from fairseq2.data.text.text_tokenizer import TextTokenEncoder
+from seamless_communication.inference import Translator
+import spacy
+import re
+from datetime import datetime
+
+router = APIRouter()
+
+class TranslateInput(BaseModel):
+    inputs: list[str]
+    model: str
+    src_lang: str
+    dst_lang: str
+
+
+class TranslateOutput(BaseModel):
+    src_lang: str
+    dst_lang: str
+    translations: Optional[list[str]] = None
+    error: Optional[str] = None
+
+
+@router.post('/t2tt')
+def t2tt(inputs: TranslateInput) -> TranslateOutput:
+    start_time = datetime.now()
+    fn = t2tt_mapping.get(inputs.model)
+    if not fn:
+        return TranslateOutput(
+            src_lang=inputs.src_lang,
+            dst_lang=inputs.dst_lang,
+            error=f'No sentence embeddings model found for {inputs.model}'
+        )
+    
+    try:
+        translations = fn(**inputs.dict())
+        log({
+            "task": "sentence_embeddings",
+            "model": inputs.model,
+            "start_time": start_time.isoformat(),
+            "time_taken": (datetime.now() - start_time).total_seconds(),
+            "inputs": inputs.inputs,
+            "outputs": translations,
+            "parameters": {
+                "src_lang": inputs.src_lang,
+                "dst_lang": inputs.dst_lang,
+            },
+        })
+        loaded_models_last_updated[inputs.model] = datetime.now()
+        return TranslateOutput(**translations)
+    except Exception as e:
+        return TranslateOutput(
+            src_lang=inputs.src_lang,
+            dst_lang=inputs.dst_lang,
+            error=str(e)
+        )
+
+cmn_nlp = spacy.load("zh_core_web_sm")
+xx_nlp = spacy.load("xx_sent_ud_sm")
+unk_re = re.compile(r"\s?<unk>|\s?⁇")
+
+def seamless_t2tt(inputs: list[str], src_lang: str, dst_lang: str = 'eng'):
+    if TEST_MODE:
+        return {
+            "src_lang": src_lang,
+            "dst_lang": dst_lang,
+            "translations": None,
+            "error": None
+        }
+
+    # Load model
+    if 'facebook/seamless-m4t-v2-large' in loaded_models:
+        translator = loaded_models['facebook/seamless-m4t-v2-large']
+    else:
+        translator = Translator(
+            model_name_or_card="seamlessM4T_v2_large",
+            vocoder_name_or_card="vocoder_v2",
+            device=device,
+            dtype=dtype,
+            apply_mintox=False,
+        )
+        loaded_models['facebook/seamless-m4t-v2-large'] = translator
+
+
+    def sent_tokenize(text, lang) -> list[str]:
+        if lang == 'cmn':
+            return [str(t) for t in cmn_nlp(text).sents]
+        return [str(t) for t in xx_nlp(text).sents]
+
+
+    def tokenize_and_translate(token_encoder: TextTokenEncoder, text: str, src_lang: str, dst_lang: str) -> str:
+        # Convert text into paragraphs and replace new lines with spaces
+        lines = [sent_tokenize(line.replace("\n", " "), src_lang) for line in text.split('\n\n') if line]
+        lines = [item for sublist in lines for item in sublist if item]
+
+        # Tokenize and translate
+        input_tokens = translator.collate([token_encoder(line) for line in lines])
+        translations = [
+                unk_re.sub("", str(t))
+                for t in translator.predict(
+                    input=input_tokens,
+                    task_str="T2TT",
+                    src_lang=src_lang,
+                    tgt_lang=dst_lang,
+                )[0]
+            ]
+        return " ".join(translations)
+
+    translations = None
+    token_encoder = translator.text_tokenizer.create_encoder(
+        task="translation", lang=src_lang, mode="source", device=translator.device
+    )
+    try:
+        translations = [tokenize_and_translate(token_encoder, text, src_lang, dst_lang) for text in inputs]
+    except Exception as e:
+        print(f"Error translating text: {e}")
+
+    return {
+        "src_lang": src_lang,
+        "dst_lang": dst_lang,
+        "translations": translations,
+        "error": None if translations else "Failed to translate text"
+    }
+
+
+# Polling every X minutes to 
+loaded_models = {}
+loaded_models_last_updated = {}
+
+t2tt_mapping = {
+    'facebook/seamless-m4t-v2-large': seamless_t2tt,
+}