57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
"""
|
|
NLLB-200 translation API (CPU, CTranslate2 int8). Part of the CRM chat pod.
|
|
Runs on port 8000 alongside Ollama/Qwen (port 11434) on the same pod.
|
|
|
|
POST /translate {"text": "...", "source": "eng_Latn", "target": "hin_Deva"}
|
|
-> {"translation": "..."}
|
|
GET /health -> {"ok": true}
|
|
|
|
Optional auth: set env CHAT_POD_KEY, then callers must send
|
|
`Authorization: Bearer <key>`.
|
|
"""
|
|
|
|
import os
|
|
|
|
import ctranslate2
|
|
import transformers
|
|
from fastapi import FastAPI, Header, HTTPException
|
|
from pydantic import BaseModel
|
|
|
|
CT2 = os.environ.get("NLLB_CT2", "/workspace/nllb-ct2")
|
|
HF = os.environ.get("NLLB_HF", "facebook/nllb-200-1.3B")
|
|
API_KEY = os.environ.get("CHAT_POD_KEY", "")
|
|
|
|
_translator = ctranslate2.Translator(CT2, device="cpu", compute_type="int8")
|
|
_tok = transformers.AutoTokenizer.from_pretrained(HF)
|
|
|
|
app = FastAPI(title="CRM Translate (NLLB-200)")
|
|
|
|
|
|
class Req(BaseModel):
|
|
text: str
|
|
source: str # FLORES code, e.g. "eng_Latn"
|
|
target: str # FLORES code, e.g. "hin_Deva"
|
|
|
|
|
|
def _check(authorization: str) -> None:
|
|
if API_KEY and authorization != f"Bearer {API_KEY}":
|
|
raise HTTPException(status_code=401, detail="unauthorized")
|
|
|
|
|
|
@app.post("/translate")
|
|
def translate(r: Req, authorization: str = Header(default="")):
|
|
_check(authorization)
|
|
text = (r.text or "").strip()
|
|
if not text:
|
|
return {"translation": ""}
|
|
_tok.src_lang = r.source
|
|
source = _tok.convert_ids_to_tokens(_tok.encode(text))
|
|
results = _translator.translate_batch([source], target_prefix=[[r.target]])
|
|
hyp = results[0].hypotheses[0][1:] # drop the leading target-language token
|
|
return {"translation": _tok.decode(_tok.convert_tokens_to_ids(hyp))}
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"ok": True}
|