Adolfo Reyna
119 lines
4.5 KiB
Python
119 lines
4.5 KiB
Python
import json
|
|
import os
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
|
|
from langdetect import DetectorFactory, LangDetectException, detect
|
|
from transformers import MarianMTModel, MarianTokenizer
|
|
|
|
DetectorFactory.seed = 0
|
|
|
|
HOST = os.getenv("MARIAN_HOST", "127.0.0.1")
|
|
PORT = int(os.getenv("MARIAN_PORT", "8000"))
|
|
MAX_INPUT_LENGTH = int(os.getenv("MARIAN_MAX_INPUT_LENGTH", "1000"))
|
|
DEFAULT_SOURCE_LANGUAGE = os.getenv("MARIAN_DEFAULT_SOURCE_LANGUAGE", "en")
|
|
SUPPORTED_LANGUAGES = {"en", "es", "fr", "da", "ar"}
|
|
MODEL_CACHE = {}
|
|
|
|
|
|
def normalize_language(value):
|
|
language = str(value or "").strip().lower().split(",")[0].split("-")[0]
|
|
return language
|
|
|
|
|
|
def detect_source_language(text):
|
|
try:
|
|
detected = normalize_language(detect(text))
|
|
if detected in SUPPORTED_LANGUAGES:
|
|
return detected
|
|
except LangDetectException:
|
|
pass
|
|
return DEFAULT_SOURCE_LANGUAGE
|
|
|
|
|
|
def get_model(source, target):
|
|
model_name = f"Helsinki-NLP/opus-mt-{source}-{target}"
|
|
if model_name not in MODEL_CACHE:
|
|
MODEL_CACHE[model_name] = (
|
|
MarianTokenizer.from_pretrained(model_name),
|
|
MarianMTModel.from_pretrained(model_name),
|
|
)
|
|
return model_name, MODEL_CACHE[model_name]
|
|
|
|
|
|
def translate_once(text, source, target):
|
|
model_name, (tokenizer, model) = get_model(source, target)
|
|
encoded = tokenizer([text], return_tensors="pt", truncation=True)
|
|
generated = model.generate(**encoded)
|
|
return tokenizer.batch_decode(generated, skip_special_tokens=True)[0], model_name
|
|
|
|
|
|
def translate(text, source, target):
|
|
if source == "auto":
|
|
source = detect_source_language(text)
|
|
if source not in SUPPORTED_LANGUAGES or target not in SUPPORTED_LANGUAGES:
|
|
raise ValueError("Only en, es, fr, da, and ar are supported")
|
|
if source == target:
|
|
return text, source, "none"
|
|
if source == "en" or target == "en":
|
|
translated, model_name = translate_once(text, source, target)
|
|
return translated, source, model_name
|
|
|
|
english, first_model = translate_once(text, source, "en")
|
|
translated, second_model = translate_once(english, "en", target)
|
|
return translated, source, f"{first_model},{second_model}"
|
|
|
|
|
|
class TranslationHandler(BaseHTTPRequestHandler):
|
|
def send_json(self, status, body):
|
|
payload = json.dumps(body).encode("utf-8")
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(payload)))
|
|
self.end_headers()
|
|
self.wfile.write(payload)
|
|
|
|
def do_GET(self):
|
|
if self.path != "/health":
|
|
self.send_json(404, {"status": "not found"})
|
|
return
|
|
self.send_json(200, {"status": "ok", "provider": "marianmt", "loadedModels": list(MODEL_CACHE)})
|
|
|
|
def do_POST(self):
|
|
if self.path != "/translate":
|
|
self.send_json(404, {"status": "not found"})
|
|
return
|
|
try:
|
|
content_length = int(self.headers.get("Content-Length", "0"))
|
|
body = json.loads(self.rfile.read(content_length).decode("utf-8"))
|
|
text = str(body.get("text") or "").strip()
|
|
source = normalize_language(body.get("sourceLang")) or "auto"
|
|
target = normalize_language(body.get("targetLang"))
|
|
if not text or not target:
|
|
self.send_json(400, {"status": "text and targetLang are required"})
|
|
return
|
|
if len(text) > MAX_INPUT_LENGTH:
|
|
self.send_json(400, {"status": f"text exceeds {MAX_INPUT_LENGTH} characters"})
|
|
return
|
|
translated, detected_source, model_name = translate(text, source, target)
|
|
self.send_json(200, {
|
|
"status": "ok",
|
|
"translatedText": translated,
|
|
"sourceLang": detected_source,
|
|
"targetLang": target,
|
|
"provider": "marianmt",
|
|
"model": model_name,
|
|
})
|
|
except (ValueError, json.JSONDecodeError) as error:
|
|
self.send_json(400, {"status": str(error)})
|
|
except Exception as error:
|
|
print(f"Translation failed: {error}", flush=True)
|
|
self.send_json(502, {"status": "Translation failed"})
|
|
|
|
def log_message(self, format_string, *args):
|
|
print(f"[marianmt] {self.address_string()} {format_string % args}", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print(f"MarianMT translation service listening on {HOST}:{PORT}", flush=True)
|
|
ThreadingHTTPServer((HOST, PORT), TranslationHandler).serve_forever()
|