def __init__(self,
file: type(defaults.file) = defaults.file,
index_name: type(defaults.index_name) = defaults.index_name,
id_col: type(defaults.id_col) = defaults.id_col,
host: type(defaults.host) = defaults.host,
port: type(defaults.port) = defaults.port,
delim: type(defaults.delim) = defaults.delim,
shards: type(defaults.shards) = defaults.shards,
verbose: type(defaults.verbose) = defaults.verbose,
**_):
"""
:param name: name of the index
:param id_col: column number of the id
:param field_col: column number of the field data
:param field_name: name of the field
:param host: host of the search api server
:param port: port the the server
:param shards: number of shards for the index
"""
self.file = file
self.index_name = index_name
self.id_col = id_col
self.host = host
self.port = port
self.delim = delim
self.shards = shards
self.logger = set_logger(self.__class__.__name__, verbose=verbose)
def __init__(self,
file: str,
name: str = 'nboost',
id_col: int = 0,
field_col: int = 1,
field_name: str = 'passage',
host: str = '0.0.0.0',
port: int = 9200,
delim: str = '\t',
shards: int = 3,
verbose: bool = False,
**_):
"""
:param name: name of the index
:param id_col: column number of the id
:param field_col: column number of the field data
:param field_name: name of the field
:param host: host of the search api server
:param port: port the the server
:param shards: number of shards for the index
"""
self.file = file
self.name = name
self.id_col = id_col
self.field_col = field_col
self.field_name = field_name
self.host = host
self.port = port
self.delim = delim
self.shards = shards
self.logger = set_logger(self.__class__.__name__, verbose=verbose)
def __init__(self,
model_dir: str = 'nboost/pt-tinybert-msmarco',
verbose: bool = defaults.verbose,
max_seq_len: int = defaults.max_seq_len,
**kwargs):
super().__init__(**kwargs)
self.logger = set_logger(model_dir, verbose=verbose)
self.max_seq_len = max_seq_len
self.logger.info('Loading from checkpoint %s' % model_dir)
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
if self.device == torch.device("cpu"):
self.logger.info("RUNNING ON CPU")
else:
self.logger.info("RUNNING ON CUDA")
torch.cuda.synchronize(self.device)
self.rerank_model = AutoModelForSequenceClassification.from_pretrained(
model_dir)
self.tokenizer = AutoTokenizer.from_pretrained(model_dir,
use_fast=True)
self.rerank_model.to(self.device, non_blocking=True)
def resolve_model(data_dir: Path, model_dir: str, model_cls: str, **kwargs):
"""Dynamically import class from a module in the CLASS_MAP. This is used
to manage dependencies within nboost. For example, you don't necessarily
want to import pytorch models everytime you boot up tensorflow..."""
logger = set_logger('resolve_model')
data_dir.mkdir(parents=True, exist_ok=True)
if 'http' in model_dir:
module = MODULE_MAP[model_cls]
model = import_class(module, model_cls)
return model(model_dir=model_dir)
model_dir = data_dir.joinpath(model_dir).absolute()
if model_dir.exists():
logger.info('Using model cache from %s', model_dir)
if model_dir.name in CLASS_MAP:
model_cls = CLASS_MAP[model_dir.name]
elif model_cls not in MODULE_MAP:
raise ImportError('Class "%s" not in %s.' % CLASS_MAP.keys())
module = MODULE_MAP[model_cls]
model = import_class(module, model_cls) # type: Type[ModelPlugin]
return model(model_dir=str(model_dir), **kwargs)
else:
if model_dir.name in CLASS_MAP:
model_cls = CLASS_MAP[model_dir.name]
module = MODULE_MAP[model_cls]
if model_dir.name in URL_MAP: # DOWNLOAD AND CACHE
url = URL_MAP[model_dir.name]
binary_path = data_dir.joinpath(Path(url).name)
if binary_path.exists():
logger.info('Found model cache in %s', binary_path)
else:
logger.info('Downloading "%s" model.', model_dir)
download_file(url, binary_path)
if binary_path.suffixes == ['.tar', '.gz']:
logger.info('Extracting "%s" from %s', model_dir,
binary_path)
extract_tar_gz(binary_path, data_dir)
else: # pass along to plugin maybe it can resolve it
model_dir = model_dir.name
model = import_class(module, model_cls) # type: Type[ModelPlugin]
return model(model_dir=str(model_dir), **kwargs)
else:
if model_cls in MODULE_MAP:
module = MODULE_MAP[model_cls]
model = import_class(module,
model_cls) # type: Type[ModelPlugin]
return model(model_dir=model_dir.name, **kwargs)
else:
raise ImportError('model_dir %s not found in %s. You must '
'set --model class to continue.' %
(model_dir.name, CLASS_MAP.keys()))
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.logger = set_logger(self.model_dir, verbose=True)
self.logger.info('Distil Loading from checkpoint %s' % self.model_dir)
self.model = DistilBertForQuestionAnswering.from_pretrained(self.model_dir)
self.tokenizer = DistilBertTokenizer.from_pretrained(self.model_dir)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(self.device)
def __init__(self,
max_query_length: type(
defaults.max_query_length) = defaults.max_query_length,
model_dir: str = defaults.qa_model_dir,
max_seq_len: int = defaults.max_seq_len,
**kwargs):
super().__init__(**kwargs)
self.model_dir = model_dir
self.max_query_length = max_query_length
self.max_seq_len = max_seq_len
self.logger = set_logger('qamodel', verbose=True)
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.logger = set_logger(self.model_dir, verbose=True)
self.device_id = 0 if torch.cuda.is_available() else -1
model_path = (Path(self.model_dir) / "model.tar.gz").absolute()
self.model = TransformerQAPredictor.from_archive(
load_archive(str(model_path), self.device_id),
predictor_name="transformer_qa")
self.logger.info('Loading AllenNLP from %s' % self.model_dir)
def __init__(self, model_dir: str, lr: float = 10e-3, batch_size: int = 4,
max_seq_len: int = 128, verbose: bool = False, filter_results: bool = False, **_):
"""Model dir will be a full path if the binary is present, and will
be just the name of the "model_dir" if it is not."""
super().__init__()
self.filter_results = filter_results
self.model_dir = model_dir
self.lr = lr
self.max_seq_len = max_seq_len
self.batch_size = batch_size
self.logger = set_logger(model_dir, verbose=verbose)
def __init__(self,
host: str = '0.0.0.0',
port: int = 8000,
backlog: int = 100,
workers: int = 10,
**kwargs):
super().__init__()
self.address = (host, port)
self.backlog = backlog
self.workers = workers
self.is_ready = Event()
self.sock = self.set_socket()
self.logger = set_logger(self.__class__.__name__)
def __init__(self, host: type(defaults.host) = defaults.host,
port: type(defaults.port) = defaults.port,
backlog: type(defaults.backlog) = defaults.backlog,
workers: type(defaults.workers) = defaults.workers,
verbose: type(defaults.verbose) = defaults.verbose,
**kwargs):
super().__init__()
self.address = (socket.gethostbyname(host), port)
self.backlog = backlog
self.workers = workers
self.is_ready = Event()
self.sock = self.set_socket()
self.logger = set_logger(self.__class__.__name__, verbose=verbose)
def __init__(self,
model_dir: type(defaults.model_dir) = defaults.model_dir,
batch_size: type(defaults.batch_size) = defaults.batch_size,
max_seq_len: type(
defaults.max_seq_len) = defaults.max_seq_len,
verbose: type(defaults.verbose) = defaults.verbose,
lr: type(defaults.lr) = defaults.lr,
**_):
"""Model dir will be a full path if the binary is present, and will
be just the name of the "model_dir" if it is not."""
super().__init__()
self.model_dir = model_dir
self.lr = lr
self.max_seq_len = max_seq_len
self.batch_size = batch_size
self.logger = set_logger(model_dir, verbose=verbose)
def __init__(self,
model: Type[BaseModel],
codex: Type[BaseCodex],
uhost: str = '0.0.0.0',
uport: int = 9200,
bufsize: int = 2048,
verbose: bool = False,
**kwargs):
super().__init__(**kwargs)
self.kwargs = kwargs
self.uaddress = (uhost, uport)
self.bufsize = bufsize
self.logger = set_logger(model.__name__, verbose=verbose)
# pass command line arguments to instantiate each component
self.model = model(verbose=verbose, **kwargs)
self.codex = codex(verbose=verbose, **kwargs)
def __init__(self,
lr: float = 10e-3,
model_dir: str = 'bert-base-uncased-msmarco',
data_dir: Path = PKG_PATH.joinpath('.cache'),
max_seq_len: int = 128,
batch_size: int = 4,
**_):
super().__init__()
self.lr = lr
self.max_seq_len = max_seq_len
self.batch_size = batch_size
self.data_dir = data_dir
if not os.path.exists(model_dir):
self.model_dir = data_dir.joinpath(model_dir).absolute()
else:
self.model_dir = Path(model_dir)
self.logger = set_logger(model_dir)
def __init__(self,
data_dir: Path = PKG_PATH.joinpath('.cache'),
model_dir: str = 'pt-bert-base-uncased-msmarco',
qa_model_dir: str = 'distilbert-base-uncased-distilled-squad',
qa_model: str = str(),
model: str = str(),
qa: bool = False,
config: str = 'elasticsearch',
verbose: bool = False,
**kwargs):
super().__init__(**kwargs)
self.qa = qa
self.data_dir = data_dir
self.data_dir.mkdir(parents=True, exist_ok=True)
self.model_dir = data_dir.joinpath(model_dir).absolute()
self.model = self.resolve_model(self.model_dir,
model,
verbose=verbose,
**kwargs)
self.logger = set_logger(self.model.__class__.__name__,
verbose=verbose)
if qa:
self.qa_model_dir = data_dir.joinpath(qa_model_dir).absolute()
self.qa_model = self.resolve_model(self.qa_model_dir, qa_model,
**kwargs)
# these are global parameters that are overrided by nboost json key
self.config = {
'model': self.model.__class__.__name__,
'model_dir': model_dir,
'qa_model': self.qa_model.__class__.__name__ if qa else None,
'qa_model_dir': qa_model_dir if qa else None,
'data_dir': str(data_dir),
**CONFIG_MAP[config],
**kwargs
}
def __init__(
self,
host: type(defaults.host) = defaults.host,
port: type(defaults.port) = defaults.port,
verbose: type(defaults.verbose) = defaults.verbose,
data_dir: type(defaults.data_dir) = defaults.data_dir,
no_rerank: type(defaults.no_rerank) = defaults.no_rerank,
model: type(defaults.model) = defaults.model,
model_dir: type(defaults.model_dir) = defaults.model_dir,
qa: type(defaults.qa) = defaults.qa,
qa_model: type(defaults.qa_model) = defaults.qa_model,
qa_model_dir: type(defaults.qa_model_dir) = defaults.qa_model_dir,
search_route: type(defaults.search_route) = defaults.search_route,
frontend_route: type(
defaults.frontend_route) = defaults.frontend_route,
status_route: type(defaults.status_route) = defaults.status_route,
debug: type(defaults.debug) = defaults.debug,
prerank: type(defaults.prerank) = defaults.prerank,
**cli_args):
self.logger = set_logger(self.__class__.__name__, verbose=verbose)
BackwardsCompatibility().set()
db = Database()
plugins = [] # type: List[Plugin]
if prerank:
preRankPlugin = PrerankPlugin()
plugins.append(preRankPlugin)
if not no_rerank:
rerank_model_plugin = resolve_model(
data_dir=data_dir,
model_dir=model_dir,
model_cls=model,
**cli_args) # type: RerankModelPlugin
plugins.append(rerank_model_plugin)
if qa:
qa_model_plugin = resolve_model(data_dir=data_dir,
model_dir=qa_model_dir,
model_cls=qa_model,
**cli_args) # type: QAModelPlugin
plugins.append(qa_model_plugin)
if debug:
debug_plugin = DebugPlugin(**cli_args)
plugins.append(debug_plugin)
static_dir = str(PKG_PATH.joinpath('resources/frontend'))
flask_app = Flask(__name__)
@flask_app.route(frontend_route, methods=['GET'])
def frontend_root():
return send_from_directory(static_dir, 'index.html')
@flask_app.route(frontend_route + '/<path:path>', methods=['GET'])
def frontend_path(path):
return send_from_directory(static_dir, path)
@flask_app.route(frontend_route + status_route)
def status_path():
configs = {}
for plugin in plugins:
configs.update(plugin.configs)
stats = db.get_stats()
return jsonify({**configs, **stats})
flask_app.url_map.add(Rule('/<path:path>', endpoint='proxy'))
@flask_app.route('/', defaults={'path': ''})
@flask_app.endpoint('proxy')
def proxy_through(path):
# parse the client request
dict_request = flask_request_to_dict_request(
flask_request) # takes the json
"""Search request."""
db_row = db.new_row()
# combine command line args and runtime args sent by request
query_args = {}
for key in list(dict_request['url']['query']):
if key in defaults.__dict__:
query_args[key] = dict_request['url']['query'].pop(key)
json_args = dict_request['body'].pop('nboost', {})
args = {**cli_args, **json_args, **query_args}
request = RequestDelegate(dict_request, **args)
request.dict['headers'].pop('Host', '')
request.set_path('url.headers.host',
'%s:%s' % (request.uhost, request.uport))
request.set_path('url.netloc',
'%s:%s' % (request.uhost, request.uport))
request.set_path('url.scheme', 'https' if request.ussl else 'http')
for plugin in plugins: # type: Plugin
plugin.on_request(request, db_row)
# get response from upstream server
start_time = perf_counter()
requests_response = dict_request_to_requests_response(dict_request)
db_row.response_time = perf_counter() - start_time
try:
dict_response = requests_response_to_dict_response(
requests_response)
except JSONDecodeError:
print(requests_response.content)
return requests_response.content
response = ResponseDelegate(dict_response, request)
response.set_path('body.nboost', {})
db_row.choices = len(response.choices)
for plugin in plugins: # type: Plugin
plugin.on_response(response, db_row)
# save stats to sql lite
db.insert(db_row)
return dict_response_to_flask_response(dict_response)
@flask_app.errorhandler(Exception)
def handle_json_response(error):
self.logger.error('', exc_info=True)
return jsonify({
'type': error.__class__.__name__,
'doc': error.__class__.__doc__,
'msg': str(error.args)
}), 500
self.run = lambda: (self.logger.critical('LISTENING %s:%s' % (
host, port)) or flask_app.run(host=host, port=port))
def execute(command: str):
"""Execute command in subprocess"""
logger = set_logger('RELEASE')
logger.info(command)
subprocess.call(command, shell=True)
