74 lines
1.8 KiB
C++
74 lines
1.8 KiB
C++
#include <httplib.h>
|
|
#include "fasttext.h"
|
|
#include <string>
|
|
#include <iostream>
|
|
#include <sstream>
|
|
#include <csignal>
|
|
|
|
httplib::Server *g_svr = nullptr;
|
|
|
|
std::string stripLabel(const std::string &label) {
|
|
const std::string prefix = "__label__";
|
|
if (label.substr(0, prefix.length()) == prefix) {
|
|
return label.substr(prefix.length());
|
|
}
|
|
return label;
|
|
}
|
|
|
|
int main(int argc, char *argv[]) {
|
|
if (argc != 3) {
|
|
std::cerr << "Usage: " << argv[0] << " <model_path> <port>" << std::endl;
|
|
return 1;
|
|
}
|
|
|
|
fasttext::FastText ft;
|
|
ft.loadModel(argv[1]);
|
|
|
|
httplib::Server svr;
|
|
|
|
svr.Get("/detect", [&ft](const httplib::Request &req, httplib::Response &res) {
|
|
if (!req.has_param("text")) {
|
|
res.status = 400;
|
|
res.set_content("Missing text parameter", "text/plain");
|
|
return;
|
|
}
|
|
|
|
auto text = req.get_param_value("text");
|
|
std::istringstream iss(text);
|
|
std::vector<std::pair<fasttext::real, std::string>> predictions;
|
|
ft.predictLine(iss, predictions, 3, 0.0);
|
|
|
|
if (predictions.empty()) {
|
|
res.status = 500;
|
|
res.set_content("Prediction failed", "text/plain");
|
|
return;
|
|
}
|
|
|
|
std::string result = "[";
|
|
|
|
bool first = true;
|
|
for (const auto& prediction : predictions) {
|
|
std::string predictionResult = "{\"language\":\"" + stripLabel(prediction.second) +
|
|
"\",\"probability\":" + std::to_string(prediction.first) + "}";
|
|
if (!first) {
|
|
result += ",";
|
|
}
|
|
result += predictionResult;
|
|
first = false;
|
|
}
|
|
|
|
result += "]";
|
|
|
|
res.set_content(result, "application/json");
|
|
});
|
|
|
|
g_svr = &svr;
|
|
std::signal(SIGINT, [](int) { if (g_svr) g_svr->stop(); });
|
|
std::signal(SIGTERM, [](int) { if (g_svr) g_svr->stop(); });
|
|
|
|
auto port = std::stoi(argv[2]);
|
|
std::cout << "Starting server on port " << port << std::endl;
|
|
svr.listen("0.0.0.0", port);
|
|
std::cout << "Server stopped" << std::endl;
|
|
return 0;
|
|
} |