def test_travel_tutorial(self):
subprocess.call('docker pull elasticsearch:7.4.2', shell=True)
subprocess.call(
'docker run -d -p 9200:9200 -p 9300:9300 '
'-e "discovery.type=single-node" elasticsearch:7.4.2',
shell=True)
for _ in range(5):
Elasticsearch().index(index='example', body={'field': 'value'})
proxy = Proxy(model_dir='shuffle-model',
model='ShuffleRerankPlugin',
uport=9200,
debug=True,
verbose=True,
query_prep='lambda query: query.split(":")[-1]')
t = Thread(target=proxy.run)
t.start()
time.sleep(2)
# search
proxy_res = requests.get('http://localhost:8000/example/_search',
params={
'q': 'field:value',
'size': 3,
'topn': 20
})
self.assertTrue(proxy_res.ok)
pprint(proxy_res.json())
response = proxy_res.json()['nboost']
self.assertEqual('value', response['query'])
self.assertEqual(3, response['topk'])
self.assertEqual(20, response['topn'])
self.assertEqual(3, len(response['cvalues']))
# fallback
# fallback_res = requests.get('http://localhost:8000/')
# self.assertTrue(fallback_res.ok)
# print(fallback_res.content.decode())
# status
status_res = requests.get('http://localhost:8000/nboost/status')
self.assertTrue(status_res.ok)
print(status_res.content.decode())
# invalid host
invalid_res = requests.get('http://localhost:8000/example/_search',
params={'uport': 2000})
print(invalid_res.content)
self.assertFalse(invalid_res.ok)
t.join()
def main(argv: List[str] = None):
name = termcolor.colored('NBoost Benchmarker', 'cyan', attrs=['underline'])
parser = ArgumentParser(
description='This is the %s for nboost, a search-api-boosting platform'
'Use the cli as normal, but instead of hosting a proxy,'
'the cli will benchmark speed and efficacy on a specified '
'dataset.' % name,
formatter_class=ArgumentDefaultsHelpFormatter)
add_default_args(parser)
parser.add_argument('--rows',
type=int,
default=-1,
help='number of rows for benchmarking')
parser.add_argument('--dataset', type=str, required=True, choices=DATASETS)
parser.add_argument('--topk', type=int, default=10)
parser.add_argument('--url', default=None, type=str)
parser.add_argument('--shards', default=1, type=str)
args = parser.parse_args(argv)
if args.dataset == 'msmarco':
args.field = 'passage'
proxy = Proxy(**vars(args))
proxy.start()
# execute dataset dependencies
benchmarker = getattr(api, args.dataset)(args)
benchmarker.benchmark()
print(requests.get('http://localhost:8000/nboost').text)
print()
proxy.close()
print('DONE')
def main():
"""Entrypoint for nboost"""
parser = set_parser()
args = parser.parse_args()
proxy = Proxy(**vars(args))
try:
proxy.start()
except KeyboardInterrupt:
proxy.close()
def test_proxy(self):
server = TestServer(port=9500, verbose=True)
proxy = Proxy(port=8000, model=ShuffleModel, codex=TestCodex,
uport=9500, verbose=True)
proxy.start()
server.start()
proxy.is_ready.wait()
server.is_ready.wait()
# search
params = dict(q='test_field;test query', topk=3)
proxy_res = requests.get('http://localhost:8000/test', params=params)
print(proxy_res.content)
self.assertTrue(proxy_res.ok)
self.assertEqual(3, len(proxy_res.json()))
server_res = requests.get('http://localhost:9500/test', params=params)
print(server_res.content)
self.assertTrue(server_res.ok)
# benchmark
params['nboost'] = '2,23'
bench_res = requests.get('http://localhost:8000/test', params=params)
print(bench_res.content)
self.assertTrue(bench_res.ok)
# fallback
fallback_res = requests.post('http://localhost:8000/only_on_server',
data=b'hello there my friend')
print('fallback:', fallback_res.content)
self.assertTrue(fallback_res.ok)
# status
status_res = requests.get('http://localhost:8000/nboost/status')
self.assertTrue(status_res.ok)
self.assertEqual(0.5, status_res.json()['vars']['upstream_mrr']['avg'])
print(status_res.content.decode())
# invalid host
proxy.uaddress = ('0.0.0.0', 2000)
invalid_res = requests.get('http://localhost:8000')
print(invalid_res.content)
self.assertFalse(invalid_res.ok)
proxy.close()
server.close()
def test_proxy(self):
server = TestServer(port=9500, verbose=True)
proxy = Proxy(host='0.0.0.0',
port=8000,
uhost='0.0.0.0',
model_dir='shuffle-model',
uport=9500,
bufsize=2048,
delim='. ',
multiplier=5,
verbose=True)
proxy.start()
server.start()
proxy.is_ready.wait()
server.is_ready.wait()
# search
params = dict(q='test_field;test query', size=3)
proxy_res = requests.get('http://localhost:8000/test/_search',
params=params)
print(proxy_res.content)
self.assertTrue(proxy_res.ok)
json = proxy_res.json()
self.assertEqual(3, len(json['hits']['hits']))
# fallback
server_res = requests.get('http://localhost:9500/test', params=params)
print(server_res.content)
self.assertTrue(server_res.ok)
# status
status_res = requests.get('http://localhost:8000/nboost/status')
self.assertTrue(status_res.ok)
print(status_res.content.decode())
# self.assertEqual(0.5, status_res.json()['vars']['upstream_mrr']['avg'])
# invalid host
proxy.config['uport'] = 2000
invalid_res = requests.get('http://localhost:8000')
print(invalid_res.content)
self.assertFalse(invalid_res.ok)
proxy.close()
server.close()
def test_proxy(self):
server = TestServer(port=9500, verbose=True)
proxy = Proxy(model_dir='shuffle-model', model='ShuffleModel', uport=9500,
verbose=True, query_prep='lambda query: query.split(";")[-1]')
proxy.start()
server.start()
proxy.is_ready.wait()
server.is_ready.wait()
# search
proxy_res = requests.get(
'http://localhost:8000/test/_search',
params={
'q': 'test_field;test query',
'size': 3,
'debug': True,
'topn': 20
}
)
self.assertTrue(proxy_res.ok)
pprint(proxy_res.json())
session = proxy_res.json()['nboost']
self.assertEqual('test query', session['query'])
self.assertEqual(3, session['topk'])
self.assertEqual(20, session['topn'])
self.assertEqual(3, len(session['cvalues']))
# fallback
server_res = requests.get('http://localhost:9500/test')
print(server_res.content)
self.assertTrue(server_res.ok)
# status
status_res = requests.get('http://localhost:8000/nboost/status')
self.assertTrue(status_res.ok)
print(status_res.content.decode())
# invalid host
invalid_res = requests.get('http://localhost:8000', params={'uport': 2000})
print(invalid_res.content)
self.assertFalse(invalid_res.ok)
proxy.close()
server.close()
def setUp(self):
self.proxy = Proxy(qa_model_dir='distilbert-base-uncased-distilled-squad', rerank=False,
qa_model='PtDistilBertQAModel', max_seq_length=32, qa=True)
def main():
"""Entrypoint for nboost"""
parser = set_parser()
args = parser.parse_args()
proxy = Proxy(**vars(args))
proxy.run()
def setUp(self):
self.proxy = Proxy(model_dir='tf-bert-base-uncased-msmarco')
def setUp(self):
self.proxy = Proxy(model_dir='pt-bert-base-uncased-msmarco',
filter_results=True)
def create_proxy(argv: List[str] = None):
"""Return proxy instance given a command line"""
parser = set_parser()
args = parser.parse_args(argv)
return Proxy(**vars(args))