marconolimits/NMT

1#include "NMTWrapper.h"2#include <chrono>3#include <ctranslate2/models/model.h>4#include <ctranslate2/translator.h>5#include <iostream>6#include <sentencepiece_processor.h>7 8using namespace std;9 10namespace NMT {11 12namespace {13 14// Match scripts/evaluate_nmt_fast.py: Translator(..., inter_threads=8, intra_threads=0)15// — C++ maps intra_threads -> ReplicaPoolConfig::num_threads_per_replica (0 = auto, see16//   ctranslate2::set_num_threads). inter_threads maps to parallel replicas; a single17//   sequential TCP worker benefits most from intra (multi-core matmul), not multiple replicas.18#if defined(_M_ARM64) || defined(__aarch64__)19constexpr size_t kReplicasPerDevice = 1;20constexpr size_t kThreadsPerReplica = 1; // HoloLens / ARM: stable single-threaded compute21#elif defined(_WIN32)22// OPENBLAS_NUM_THREADS=1 in main.cpp caps BLAS; CT2 can still use OMP_NUM_THREADS for ops.23// kThreadsPerReplica=0 matches evaluate_nmt_fast.py intra_threads=0 (auto / OMP env).24constexpr size_t kReplicasPerDevice = 1;25constexpr size_t kThreadsPerReplica = 0;26#else27constexpr size_t kReplicasPerDevice = 1;28constexpr size_t kThreadsPerReplica =29    0; // Linux/macOS: 0 = intra_threads=0 (auto)30#endif31 32} // namespace33 34struct NMTWrapper::Impl {35  unique_ptr<ctranslate2::Translator> translator;36  unique_ptr<sentencepiece::SentencePieceProcessor> sp_processor;37 38  Impl(const std::string &model_path, const std::string &tokenizer_path) {39    ctranslate2::models::ModelLoader loader(model_path);40    loader.device = ctranslate2::Device::CPU;41    loader.num_replicas_per_device = kReplicasPerDevice;42 43    ctranslate2::ReplicaPoolConfig pool_config;44    pool_config.num_threads_per_replica = kThreadsPerReplica;45 46    translator =47        make_unique<ctranslate2::Translator>(loader, pool_config);48 49    // 2. Load the Tokenizer (SentencePiece)50    sp_processor = make_unique<sentencepiece::SentencePieceProcessor>();51    const auto status = sp_processor->Load(tokenizer_path);52    if (!status.ok()) {53      cerr << "❌ Failed to load SentencePiece model: " << status.ToString()54           << endl;55    }56  }57 58  string translate(const string &english_text) {59    if (!translator || !sp_processor)60      return "Error: Engine not initialized";61 62    using clock = std::chrono::steady_clock;63    auto t0 = clock::now();64 65    // --- A. TOKENIZE ---66    vector<string> tokens;67    sp_processor->Encode(english_text, &tokens);68 69    // --- THE FIX: ADD SOURCE LANGUAGE TAGS ---70    // NLLB *requires* the input to end with "</s>" and the source language71    // code. If we miss this, the model hallucinates (Holo Holo Holo...)72    tokens.push_back("</s>");73    tokens.push_back("eng_Latn");74 75    vector<vector<string>> batch_input = {tokens};76 77    auto t1 = clock::now();78 79    // --- B. SET OPTIONS (match evaluate_nmt_fast.py: beam_size=1, max_decoding_length=256) ---80    ctranslate2::TranslationOptions options;81    options.beam_size = 1;82    options.max_decoding_length = 256;83 84    // --- C. DEFINE TARGET PREFIX ---85    // This tells the model: "Start generating in Italian"86    vector<string> target_prefix = {"ita_Latn"};87    vector<vector<string>> batch_target_prefix = {target_prefix};88 89    // --- D. TRANSLATE ---90    ctranslate2::TranslationResult result = translator->translate_batch(91        batch_input, batch_target_prefix, options)[0];92 93    auto t2 = clock::now();94 95    // --- E. DETOKENIZE ---96    string italian_text;97    sp_processor->Decode(result.output(), &italian_text);98 99    auto t3 = clock::now();100 101    auto ms = [](clock::time_point a, clock::time_point b) {102      return std::chrono::duration_cast<std::chrono::milliseconds>(b - a).count();103    };104    cout << "⏱️ NMT: " << ms(t0, t3) << " ms total — tokenize " << ms(t0, t1)105         << " ms | CTranslate2 " << ms(t1, t2) << " ms | detokenize " << ms(t2, t3)106         << " ms\n";107 108    // --- F. CLEANUP (Remove the 'ita_Latn' tag) ---109    string tag = "ita_Latn";110    // Check if text starts with tag (handling potential spaces)111    if (italian_text.find(tag) == 0) {112      // Remove tag + any following space113      size_t remove_len = tag.length();114      if (italian_text.length() > remove_len &&115          italian_text[remove_len] == ' ') {116        remove_len++;117      }118      italian_text.erase(0, remove_len);119    }120 121    return italian_text;122  }123};124 125// --- PIMPL BOILERPLATE ---126NMTWrapper::NMTWrapper(const string &model_path, const string &tokenizer_path)127    : impl(make_unique<Impl>(model_path, tokenizer_path)) {}128 129NMTWrapper::~NMTWrapper() = default;130 131string NMTWrapper::translate(const string &text) {132  return impl->translate(text);133}134} // namespace NMT