Commit e2e77c5c authored by Wael Ramadan's avatar Wael Ramadan
Browse files

use_gpu env variable

parent 6bcc9962
Pipeline #6810 passed with stage
in 4 minutes and 44 seconds
......@@ -8,3 +8,4 @@ MLP_WORKER_DEFAULT_LANGUAGE_CODE = os.getenv("MLP_WORKER_DEFAULT_LANGUAGE_CODES"
MLP_WORKER_RESOURCE_DIR = os.getenv("MLP_WORKER_RESOURCE_DIR", "../data")
MLP_WORKER_BROKER = os.getenv("MLP_WORKER_REDIS_URL", "redis://mlp-redis:6379/1")
MLP_WORKER_RESULT_BACKEND = os.getenv("MLP_WORKER_REDIS_RESULT_BACKEND", MLP_WORKER_BROKER)
MLP_USE_GPU = os.getenv("MLP_USE_GPU", "true")
......@@ -4,7 +4,7 @@ from typing import List, Optional
from celery import Celery, shared_task
from mlp_docparser import MLPDocParser
from settings import (MLP_WORKER_BROKER, MLP_WORKER_DEFAULT_LANGUAGE_CODE, MLP_WORKER_LANGUAGE_CODES, MLP_WORKER_RESOURCE_DIR, MLP_WORKER_RESULT_BACKEND)
from settings import (MLP_WORKER_BROKER, MLP_WORKER_DEFAULT_LANGUAGE_CODE, MLP_WORKER_LANGUAGE_CODES, MLP_WORKER_RESOURCE_DIR, MLP_WORKER_RESULT_BACKEND, MLP_USE_GPU)
from texta_mlp.mlp import MLP
......@@ -31,7 +31,8 @@ def load_mlp():
ml_processor = MLP(
language_codes=MLP_WORKER_LANGUAGE_CODES,
default_language_code=MLP_WORKER_DEFAULT_LANGUAGE_CODE,
resource_dir=MLP_WORKER_RESOURCE_DIR
resource_dir=MLP_WORKER_RESOURCE_DIR,
use_gpu=MLP_USE_GPU
)
logging.info("Successfully loaded MLP models.")
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment