tokenizer.cpp

#include "tokenizer.hpp"
#include <algorithm>
#include <cctype>
#include <fstream>
#include <map>
#include <sstream>

Tokenizer::Tokenizer() {
    id2word_.push_back("<UNK>");
    word2id_["<UNK>"] = 0;
}

std::string Tokenizer::normalize(const std::string& text) const {
    std::string out;
    out.reserve(text.size());
    for (unsigned char c : text) {
        if (std::isalnum(c))
            out += static_cast<char>(std::tolower(c));
        else if (std::isspace(c) || c == '?' || c == '!' || c == '.' || c == ',')
            out += ' ';
    }
    return out;
}

std::vector<std::string> Tokenizer::split(const std::string& text) const {
    std::vector<std::string> tokens;
    std::istringstream ss(text);
    std::string tok;
    while (ss >> tok) tokens.push_back(tok);
    return tokens;
}

void Tokenizer::build_vocab(const std::vector<std::string>& texts) {
    std::map<std::string, int> freq;
    for (const auto& text : texts) {
        auto words = split(normalize(text));
        for (const auto& w : words) freq[w]++;
    }

    std::vector<std::pair<int, std::string>> sorted;
    for (const auto& [w, c] : freq) sorted.emplace_back(c, w);
    std::sort(sorted.rbegin(), sorted.rend());

    int limit = VOCAB_SIZE - 1; // slot 0 = <UNK>
    for (const auto& [cnt, w] : sorted) {
        if (static_cast<int>(id2word_.size()) >= VOCAB_SIZE) break;
        if (word2id_.count(w) == 0) {
            word2id_[w] = static_cast<int>(id2word_.size());
            id2word_.push_back(w);
        }
        if (--limit == 0) break;
    }
}

std::vector<int> Tokenizer::tokenize(const std::string& text) const {
    auto words = split(normalize(text));
    std::vector<int> ids;
    ids.reserve(words.size());
    for (const auto& w : words) {
        auto it = word2id_.find(w);
        ids.push_back(it == word2id_.end() ? UNK_ID : it->second);
        if (static_cast<int>(ids.size()) >= MAX_LEN) break;
    }
    return ids;
}

bool Tokenizer::save(const std::string& path) const {
    std::ofstream f(path);
    if (!f) return false;
    for (const auto& w : id2word_) f << w << "\n";
    return f.good();
}

bool Tokenizer::load(const std::string& path) {
    std::ifstream f(path);
    if (!f) return false;
    word2id_.clear();
    id2word_.clear();
    std::string line;
    while (std::getline(f, line)) {
        if (!line.empty()) {
            word2id_[line] = static_cast<int>(id2word_.size());
            id2word_.push_back(line);
        }
    }
    return !id2word_.empty();
}