139 lines
4.8 KiB
C++
139 lines
4.8 KiB
C++
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT license.
|
|
|
|
#include <ctime>
|
|
#include <functional>
|
|
#include <iomanip>
|
|
#include <string>
|
|
#include <cstdlib>
|
|
#include <codecvt>
|
|
#include <boost/program_options.hpp>
|
|
|
|
#include <restapi/server.h>
|
|
|
|
using namespace diskann;
|
|
namespace po = boost::program_options;
|
|
|
|
std::unique_ptr<Server> g_httpServer(nullptr);
|
|
std::vector<std::unique_ptr<diskann::BaseSearch>> g_inMemorySearch;
|
|
|
|
void setup(const utility::string_t &address, const std::string &typestring)
|
|
{
|
|
web::http::uri_builder uriBldr(address);
|
|
auto uri = uriBldr.to_uri();
|
|
|
|
std::cout << "Attempting to start server on " << uri.to_string() << std::endl;
|
|
|
|
g_httpServer = std::unique_ptr<Server>(new Server(uri, g_inMemorySearch, typestring));
|
|
std::cout << "Created a server object" << std::endl;
|
|
|
|
g_httpServer->open().wait();
|
|
ucout << U"Listening for requests on: " << address << std::endl;
|
|
}
|
|
|
|
void teardown(const utility::string_t &address)
|
|
{
|
|
g_httpServer->close().wait();
|
|
}
|
|
|
|
int main(int argc, char *argv[])
|
|
{
|
|
std::string data_type, index_file, data_file, address, dist_fn, tags_file;
|
|
uint32_t num_threads;
|
|
uint32_t l_search;
|
|
|
|
po::options_description desc{"Arguments"};
|
|
try
|
|
{
|
|
desc.add_options()("help,h", "Print information on arguments");
|
|
desc.add_options()("data_type", po::value<std::string>(&data_type)->required(), "data type <int8/uint8/float>");
|
|
desc.add_options()("address", po::value<std::string>(&address)->required(), "Web server address");
|
|
desc.add_options()("data_file", po::value<std::string>(&data_file)->required(),
|
|
"File containing the data found in the index");
|
|
desc.add_options()("index_path_prefix", po::value<std::string>(&index_file)->required(),
|
|
"Path prefix for saving index file components");
|
|
desc.add_options()("num_threads,T", po::value<uint32_t>(&num_threads)->required(),
|
|
"Number of threads used for building index");
|
|
desc.add_options()("l_search", po::value<uint32_t>(&l_search)->required(), "Value of L");
|
|
desc.add_options()("dist_fn", po::value<std::string>(&dist_fn)->default_value("l2"),
|
|
"distance function <l2/mips>");
|
|
desc.add_options()("tags_file", po::value<std::string>(&tags_file)->default_value(std::string()),
|
|
"Tags file location");
|
|
po::variables_map vm;
|
|
po::store(po::parse_command_line(argc, argv, desc), vm);
|
|
if (vm.count("help"))
|
|
{
|
|
std::cout << desc;
|
|
return 0;
|
|
}
|
|
po::notify(vm);
|
|
}
|
|
catch (const std::exception &ex)
|
|
{
|
|
std::cerr << ex.what() << std::endl;
|
|
return -1;
|
|
}
|
|
diskann::Metric metric;
|
|
if (dist_fn == std::string("l2"))
|
|
metric = diskann::Metric::L2;
|
|
else if (dist_fn == std::string("mips"))
|
|
metric = diskann::Metric::INNER_PRODUCT;
|
|
else
|
|
{
|
|
std::cout << "Error. Only l2 and mips distance functions are supported" << std::endl;
|
|
return -1;
|
|
}
|
|
|
|
if (data_type == std::string("float"))
|
|
{
|
|
auto searcher = std::unique_ptr<diskann::BaseSearch>(
|
|
new diskann::InMemorySearch<float>(data_file, index_file, tags_file, metric, num_threads, l_search));
|
|
g_inMemorySearch.push_back(std::move(searcher));
|
|
}
|
|
else if (data_type == std::string("int8"))
|
|
{
|
|
auto searcher = std::unique_ptr<diskann::BaseSearch>(
|
|
new diskann::InMemorySearch<int8_t>(data_file, index_file, tags_file, metric, num_threads, l_search));
|
|
g_inMemorySearch.push_back(std::move(searcher));
|
|
}
|
|
else if (data_type == std::string("uint8"))
|
|
{
|
|
auto searcher = std::unique_ptr<diskann::BaseSearch>(
|
|
new diskann::InMemorySearch<uint8_t>(data_file, index_file, tags_file, metric, num_threads, l_search));
|
|
g_inMemorySearch.push_back(std::move(searcher));
|
|
}
|
|
else
|
|
{
|
|
std::cerr << "Unsupported data type " << argv[2] << std::endl;
|
|
}
|
|
|
|
while (1)
|
|
{
|
|
try
|
|
{
|
|
setup(address, data_type);
|
|
std::cout << "Type 'exit' (case-sensitive) to exit" << std::endl;
|
|
std::string line;
|
|
std::getline(std::cin, line);
|
|
if (line == "exit")
|
|
{
|
|
teardown(address);
|
|
g_httpServer->close().wait();
|
|
exit(0);
|
|
}
|
|
}
|
|
catch (const std::exception &ex)
|
|
{
|
|
std::cerr << "Exception occurred: " << ex.what() << std::endl;
|
|
std::cerr << "Restarting HTTP server";
|
|
teardown(address);
|
|
}
|
|
catch (...)
|
|
{
|
|
std::cerr << "Unknown exception occurreed" << std::endl;
|
|
std::cerr << "Restarting HTTP server";
|
|
teardown(address);
|
|
}
|
|
}
|
|
}
|