lollms/tests/endoints_unit_tests/cpp/test_generation.cpp
Saifeddine ALOUI 4cb313c8fa test
2023-06-18 22:51:44 +02:00

96 lines
3.4 KiB
C++

#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <thread>
#include <chrono>
// Don't forget to download these libraries
// git clone https://github.com/socketio/socket.io-client-cpp.git
#include <cxxopts.hpp>
#include <sio_client.h>
// Connect to the Socket.IO server
sio::client socket_client;
// Event handler for receiving generated text
void text_generated(const sio::event& event) {
std::cout << "Generated text: " << event.get_message()->get_string() << std::endl;
}
void test_generate_text(const std::string& host, int port, const std::string& text_file) {
// Read the text file and split by multiple newlines
std::cout << "Loading file" << std::endl;
std::ifstream file(text_file);
std::string prompt;
std::vector<std::string> prompts;
while (std::getline(file, prompt, '\n')) {
if (!prompt.empty())
prompts.push_back(prompt);
}
bool is_ready = false;
// Event handler for successful connection
socket_client.set_open_listener([&]() {
std::cout << "Connected to Socket.IO server" << std::endl;
for (const auto& prompt : prompts) {
if (!prompt.empty()) {
// Trigger the 'generate_text' event with the prompt
is_ready = false;
std::cout << "Sending prompt: " << prompt << std::endl;
socket_client.socket()->emit("generate_text", sio::object_message::create_object({
{ "prompt", prompt },
{ "personality", -1 },
{ "n_predicts", 1024 }
}));
while (!is_ready)
std::this_thread::sleep_for(std::chrono::milliseconds(100));
}
}
});
socket_client.socket()->on("text_chunk", [&](sio::event& event) {
std::cout << event.get_message()->get_map()["chunk"]->get_string();
});
socket_client.socket()->on("text_generated", [&](sio::event& event) {
std::cout << "text_generated_ok" << std::endl;
std::cout << event.get_message()->get_map()["text"]->get_string() << std::endl;
is_ready = true;
});
std::cout << "Connecting to http://" << host << ":" << port << std::endl;
// Connect to the Socket.IO server
socket_client.connect(host + ":" + std::to_string(port));
// Start the event loop
socket_client.socket()->io_service()->poll();
}
int main(int argc, char** argv) {
cxxopts::Options options("Socket.IO endpoint test", "Usage: ./executable [options]");
options.add_options()
("h,host", "Socket.IO server host", cxxopts::value<std::string>()->default_value("localhost"))
("p,port", "Socket.IO server port", cxxopts::value<int>()->default_value("9600"))
("t,text-file", "Path to the text file", cxxopts::value<std::string>());
auto parsed_args = options.parse(argc, argv);
std::string host = parsed_args["host"].as<std::string>();
int port = parsed_args["port"].as<int>();
std::string text_file = parsed_args["text-file"].as<std::string>();
// Verify if the text file exists
std::ifstream file_stream(text_file);
if (!file_stream) {
std::cerr << "Error: The provided text file '" << text_file << "' does not exist." << std::endl;
return 1;
}
// Run the test with provided arguments
test_generate_text(host, port, text_file);
return 0;
}