pietre/rss2
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: remote worker GPU support with nllb-1.3B

Browse source 
- Update Dockerfile.remote-worker with CUDA support
- Update docker-compose.yml with GPU settings
- Update remote_translator_worker.py to use nllb-200-1.3B
- Update worker_ws.go with model URL handler
This commit is contained in:
lordpietre 2026-04-01 03:55:40 +00:00
parent 8093cc1e40
commit 98e0cd0e46
4 changed files with 65 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

8
workers/remote_translator_worker.py
View file

@ -5,6 +5,7 @@ import json
import logging
import re
import threading
import hashlib
from typing import List, Optional
import websocket
@ -25,7 +26,7 @@ WORKER_SERVER = os.environ.get("WORKER_SERVER", "ws://localhost:8080/ws/worker")
DEVICE = os.environ.get("CT2_DEVICE", "cpu")
MODEL_PATH = os.environ.get("CT2_MODEL_PATH", "/app/models/nllb-ct2")
COMPUTE_TYPE = os.environ.get("CT2_COMPUTE_TYPE", "int8")
UNIVERSAL_MODEL = os.environ.get("UNIVERSAL_MODEL", "facebook/nllb-200-distilled-600M")
UNIVERSAL_MODEL = os.environ.get("UNIVERSAL_MODEL", "facebook/nllb-200-1.3B")
LANG_CODE_MAP = {
"en": "eng_Latn", "es": "spa_Latn", "fr": "fra_Latn", "de": "deu_Latn",
@ -60,7 +61,8 @@ def ensure_model():
model_bin = os.path.join(MODEL_PATH, "model.bin")
if not os.path.exists(model_bin):
LOG.info(f"CTranslate2 model not found at {MODEL_PATH}, converting...")
LOG.info(f"CTranslate2 model not found at {MODEL_PATH}")
LOG.info("Downloading from HuggingFace...")
convert_model()
LOG.info(f"Loading CTranslate2 model from {MODEL_PATH} on {DEVICE}")
@ -82,6 +84,8 @@ def convert_model():
quantization = COMPUTE_TYPE if COMPUTE_TYPE != "auto" else "int8"
LOG.info(f"Converting {UNIVERSAL_MODEL} to CTranslate2 format...")
cmd = [
"ct2-transformers-converter",
"--model", UNIVERSAL_MODEL,