translator.cpp

#include "common.h"
#include "llama.h"
#include "main.hpp"
#include "sampling.h"
#include "tools/tiny_sha1.hpp"
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <deque>
#include <filesystem>
#include <map>
#include <mutex>
#include <numeric>
#include <thread>
#include <unordered_map>

extern volatile bool g_stop;

extern std::string g_neg;
extern std::string example;
extern std::string iprefix;
extern std::string isuffix;

extern std::string cache_prefix;

extern uint levenshtein_distance(std::u16string_view, std::u16string_view);

extern std::string print_line(line_id id, std::string* line, bool stream);

extern void update_segment(uint seg, bool upd_names, uint count);

namespace fs = std::filesystem;

decltype(g_replaces) make_name_replaces(std::string_view text)
{
	static const std::vector<std::pair<std::vector<std::string>, std::vector<std::string>>> s_suffix_map{
		// clang-format off
		{{"さん"}, {"-san"}},
		{{"君", "くん"}, {"-kun"}},
		{{"様", "さま", "しゃま", "ちゃま"}, {"-sama", "-shama", "-chama"}},
		{{"ちゃん", "たん"}, {"-chan", "-tan"}},
		{{"どの", "殿"}, {"-dono"}},
		{{"先輩", "せんぱい", "センパイ"}, {"-senpai", "-sempai"}},
		{{"先生", "せんせい", "センセイ"}, {"-sensei"}}
		// clang-format on
	};

	decltype(g_replaces) result;

	for (const auto& [orig_name, tr_name] : g_speakers) {
		// Filter non-names (TODO: this should be much more complicated)
		if (tr_name.empty() || tr_name.find_first_of("?& ") + 1)
			continue;
		std::string_view name = orig_name;
		name -= ":";
		std::size_t fpos = s_suffix_map.size(), pos = 0;
		std::vector<char> features(fpos + 1, 0);
		while (pos < text.size()) {
			// TODO: convert tr_name to hiragana/katakana and use for search as well
			pos = text.find(name, pos);
			if (pos + 1 == 0)
				break;
			pos += name.size();
			bool found = false;
			// Test each suffix in set
			for (std::size_t f = 0; f < s_suffix_map.size(); f++) {
				for (const auto& suf : s_suffix_map[f].first) {
					if (text.substr(pos).starts_with(suf)) {
						found = true;
						features[f] |= 1;
						fpos = f;
						break;
					}
				}
			}
			if (!found)
				features.back() |= 1;
		}

		// Conflicts or nothing found
		if (std::accumulate(features.begin(), features.end(), 0) != 1)
			continue;
		std::string from(tr_name - ":");
		std::string to(from);
		if (!features.back()) {
			to += s_suffix_map[fpos].second[0];
			result.emplace_back(from, to);
		} else {
			// Disable "remove all suffixes" for now
			// Add dummy replace to indicate the presence of the name
			result.emplace_back(from, from);
			continue;
		}
		for (std::size_t f = 0; f < s_suffix_map.size(); f++) {
			if (!features[f]) {
				for (const auto& tr_sfx : s_suffix_map[f].second) {
					result.emplace_back(from + tr_sfx, to);
				}
			}
		}
	}

	return result;
}

// Compose KV-cache filename from segment name, line pos and hashed history+tr_text
fs::path get_cache_file(const fs::path& parent, uint hpos, line_id id)
{
	std::string result = g_lines.segs[id.first].src_name;
	result += "_";
	result += std::to_string(id.second);
	result += "_";
	char buf[42]{};
	auto& tr_text = g_lines[id].tr_text;
	{
		sha1::SHA1 s;
		// Only IDs are hashed which is supposed to distinguish only "position" in the history
		s.processBytes(g_history.data(), hpos * sizeof(g_history[0]));
		s.processBytes(tr_text.data(), tr_text.size());
		s.processBytes(&hpos, 4);
		std::uint32_t digest[5];
		s.getDigest(digest);
		std::snprintf(buf, 41, "%08x%08x%08x%08x%08x", digest[0], digest[1], digest[2], digest[3], digest[4]);
	}
	result += buf;
	result += ".kvseq";
	if (result.size() > 255)
		result.erase(0, result.size() - 255); // TODO: something better
	return parent / result;
}

std::vector<uint> get_inf_cache(gpt_params& params, line_id id, std::size_t pos_max)
{
	if (pos_max >= g_history.size()) {
		// Handle underflow from subtraction
		return {};
	}

	// Map encountered pairs of characters
	static const auto s_pairs = []() {
		std::unordered_map<uint, std::vector<line_id>> pairs;
		for (uint i = 0; i < g_lines.segs.size(); i++) {
			auto& seg = g_lines.segs[i];
			for (uint j = 0; j < seg.lines.size(); j++) {
				line_id id{i, j};
				auto& sq = seg.lines[j].sq2;
				sq = squeeze2_line(seg.lines[j].text);
				if (sq.empty())
					continue;
				for (uint k = 0; k < sq.size() - 1; k++) {
					const uint ch_pair = sq[k] + (sq[k + 1] << 16);
					auto& list = pairs[ch_pair];
					if (list.empty() || list.back() != id)
						list.push_back(id);
				}
			}
		}
		return pairs;
	}();

	// Map ids to their (latest) position in history (TODO: make it part of g_history)
	std::map<line_id, std::size_t> hist_map;
	for (auto& id0 : g_history) {
		hist_map[id0] = &id0 - g_history.data();
	}

	// Vector of history positions
	std::vector<uint> result;

	// Map relevancy
	struct rel_ref {
		line_id id;
		uint dist; // lev.distance
		uint sim;  // similarity index
	};
	std::map<std::size_t, rel_ref> rel_map;

	auto& sq = g_lines[id].sq2;
	if (sq.empty())
		return {};
	for (uint k = 0; k < sq.size() - 1; k++) {
		const uint ch_pair = sq[k] + (sq[k + 1] << 16);
		for (auto id2 : s_pairs.at(ch_pair)) {
			const auto fnd = hist_map.find(id2);
			if (fnd != hist_map.end() && fnd->second <= pos_max) {
				auto& rel = rel_map[fnd->second];
				if (rel.id != id2) {
					rel.id = id2;
					auto& sq1 = g_lines[id2].sq2;
					rel.dist = levenshtein_distance(sq, sq1);
					rel.sim = std::max(sq.size(), sq1.size()) - rel.dist;
				}
			}
		}
	}
	std::vector<decltype(rel_map)::reverse_iterator> rel_list;
	for (auto it = rel_map.rbegin(); it != rel_map.rend(); it++) {
		rel_list.push_back(it);
	}
	std::stable_sort(rel_list.begin(), rel_list.end(), [](auto& a, auto& b) {
		// Sort by similarity in descending order
		return a->second.sim > b->second.sim;
	});

	uint tokens = 0;
	for (auto& it : rel_list) {
		// Cut elements with low similarity
		auto& ref = it->second;
		tokens += g_lines[ref.id].tokens;
		if (tokens > params.n_ctx / 8 + 0u || ref.sim <= 1) {
			rel_list.resize(&it - rel_list.data());
			break;
		}
	}
	std::sort(rel_list.begin(), rel_list.end(), [](auto& a, auto& b) {
		// Sort again by history order
		return a->first < b->first;
	});
	result.reserve(rel_list.size());
	for (auto& it : rel_list) {
		result.push_back(it->first);
	}
	return result;
}

bool translate(gpt_params& params, line_id id, tr_cmd cmd)
{
	static const auto s_main_tid = std::this_thread::get_id();
	static std::atomic<uint> stop_sig = -1; // stop signal for thread, new id.second is sent
	static std::atomic<uint> work_res = 0;	// number of translated lines in segment, done by worker
	static std::condition_variable work_cv;
	static std::thread worker;

	static auto is_stopped = []() -> bool {
		if (g_stop)
			return true;
		return std::this_thread::get_id() != s_main_tid && stop_sig + 1;
	};

	static auto join_worker = [](uint val) {
		if (worker.joinable()) {
			std::lock_guard{g_mutex}, stop_sig = val;
			work_cv.notify_all();
			worker.join();
		}
	};

	// Initialize llama.cpp
	static const auto init_result = [&]() -> llama_init_result {
		llama_backend_init();
		llama_numa_init(params.numa);
		return llama_init_from_gpt_params(params);
	}();

	static const auto model = init_result.model;
	static const auto ctx = init_result.context;
	static std::vector<llama_token> tokens; // Current tokens (not necessarily decoded)
	static std::deque<uint> chunks;			// Size (in tokens) of each translated message block
	static uint decoded = 0;				// Number of token successfully llama_decode()'d
	static uint segment = -1;				// Current segment
	static uint hist_pos = -1;				// Position in g_history corresponding to chunks[0]

	// Something like "$HOME/.cache/VNSleuth/GameXXX"
	static const fs::path s_cache_path = []() {
		fs::path result;
#if defined(__APPLE__)
		result = std::getenv("HOME");
		result /= "Library";
		result /= "Caches";
#elif defined(_WIN32)
		result = std::getenv("LOCALAPPDATA");
#else
		if (auto v = std::getenv("XDG_CACHE_HOME")) {
			result = v;
		} else {
			result = std::getenv("HOME");
			result /= ".cache";
		}
#endif
		result /= "VNSleuth";
		return result;
	}() / cache_prefix;

	static const struct _init_t {
		explicit operator bool() const { return model && ctx; }
		_init_t(gpt_params&)
		{
			if (!model || !ctx)
				return;
			fs::create_directories(s_cache_path);
		}
		~_init_t()
		{
			join_worker(0);
			llama_free(ctx);
			llama_free_model(model);
			llama_backend_free();
		}
	} _init{params};

	if (!_init) {
		std::cerr << "Failed to initialize llama model." << std::endl;
		return false;
	}

	if (cmd == tr_cmd::sync) {
		// Abort background worker and possibly discard work
		// Send Id to start discarding from, obviously can't use -1
		join_worker(std::min<uint>(id.second, -2));
		return true;
	}

	if (std::this_thread::get_id() == s_main_tid) {
		// Stop background worker first
		if (cmd == tr_cmd::kick && segment == id.first && worker.joinable()) {
			// Kick only once last translated line has been read
			uint tr_lines = work_res;
			uint rd_lines = id.second + 1;
			// Compare translated count with the number of read lines
			if (tr_lines > rd_lines) {
				// Only notify worker
				work_cv.notify_all();
				return true;
			} else if (tr_lines != rd_lines) {
				std::fprintf(stderr, "%sError: Kicked from untranslated line: %u<%u\n", g_esc.reset, tr_lines, rd_lines);
			}
		}

		if (cmd == tr_cmd::kick || cmd == tr_cmd::translate) {
			join_worker(id.first != segment ? 0 : id.second + 1);
		}
	}

	// Remove first message block (returns number of tokens to erase)
	static auto eject_first = [&params](uint last_count) -> uint {
		if (chunks.empty()) {
			return 0;
		}

		auto count = chunks.front();
		chunks.pop_front();
		if (decoded > 0u + params.n_keep) {
			// std::cerr << "*Used cells: " << llama_get_kv_cache_used_cells(ctx) << std::endl;
			// std::cerr << "*Decoded: " << decoded << std::endl;
			// std::cerr << "*Tokens: " << tokens.size() << std::endl;
			const auto p0 = params.n_keep;
			const auto p1 = params.n_keep + count;
			// Copy fragment to the sequence 1 and shift it to zero pos
			llama_kv_cache_seq_cp(ctx, 0, 1, p0, p1);
			llama_kv_cache_seq_add(ctx, 1, p0, p1, -p0);
			auto fname = (std::lock_guard{g_mutex}, get_cache_file(s_cache_path, hist_pos, g_history.at(hist_pos)));
			if (!fs::is_regular_file(fname)) {
				// Need to update K-shift and stuff
				llama_kv_cache_update(ctx);
				llama_state_seq_save_file(ctx, fname.c_str(), 1, tokens.data() + params.n_keep + last_count, count);
			}
			// Undo shifting cells; delete sequence 1
			llama_kv_cache_seq_add(ctx, 1, 0, count, p0);
			llama_kv_cache_seq_keep(ctx, 0);
			// Finally remove cells
			if (!llama_kv_cache_seq_rm(ctx, 0, p0, p1))
				throw std::runtime_error("llama_kv_cache_seq_rm 1");
			llama_kv_cache_seq_add(ctx, 0, p1, -1, 0u - count);
			llama_kv_cache_defrag(ctx);
			decoded -= count;
			if (llama_get_kv_cache_used_cells(ctx) + 0u != decoded) {
				std::cerr << "Used cells: " << llama_get_kv_cache_used_cells(ctx) << std::endl;
				std::cerr << "Decoded: " << decoded << std::endl;
				// llama_kv_cache_view view = llama_kv_cache_view_init(ctx, 2);
				// llama_kv_cache_view_update(ctx, &view);
				// llama_kv_cache_dump_view(view);
				throw std::runtime_error("used cells overflow");
			}
			if (decoded > 0u + params.n_ctx)
				throw std::out_of_range("ctx underflow");
			if (decoded < 0u + params.n_keep)
				throw std::out_of_range("eject_first failed");
		}
		hist_pos++;
		return count;
	};

	static const auto unload = []() {
		uint count = 0;
		while (!chunks.empty()) {
			count += eject_first(count);
		}
		tokens.clear();
		decoded = 0;
		llama_kv_cache_clear(ctx);
		hist_pos = 0;
	};

	// Remove last message block(s)
	static auto eject_bunch = [&params](uint i, bool locked = true) {
		if (i > chunks.size()) {
			unload();
			if (llama_add_bos_token(model)) {
				tokens.push_back(llama_token_bos(model));
			}

			for (auto& t : llama_tokenize(model, params.prompt + example, false)) {
				tokens.push_back(t);
			}

			params.n_keep = tokens.size();
			return;
		}
		while (i--) {
			auto count = chunks.back();
			chunks.pop_back();
			// Purge KV cache file for discarded translation (TODO: reduce fs access?)
			if (auto hpos = hist_pos + chunks.size(); locked && hpos < g_history.size()) {
				auto fname = get_cache_file(s_cache_path, hpos, g_history[hpos]);
				fs::remove(fname);
			}
			if (decoded == tokens.size()) {
				if (!llama_kv_cache_seq_rm(ctx, 0, decoded - count, -1))
					throw std::runtime_error("llama_kv_cache_seq_rm last");
				if (llama_get_kv_cache_used_cells(ctx) + 0u != tokens.size() - count)
					throw std::runtime_error("used cells in eject_bunch");
				decoded -= count;
				if (decoded > tokens.size() || decoded < 0u + params.n_keep)
					throw std::out_of_range("eject_bunch decoded=" + std::to_string(decoded));
			}
			tokens.resize(tokens.size() - count);
			g_stats->raw_discards++;
		}
	};

	// Eject old translations if necessary keeping only 7/8 of the context full (TODO)
	static auto eject_start = [&params](uint ahead = 0) -> bool {
		uint count = 0;
		while (tokens.size() - ahead - count + params.n_predict > params.n_ctx * 7 / 8 - 1u) {
			if (chunks.empty()) {
				std::cerr << "Prompt too big or context is too small" << std::endl;
				return false;
			}
			count += eject_first(count);
		}

		tokens.erase(tokens.begin() + params.n_keep, tokens.begin() + (params.n_keep + count));
		// Apply defrag
		if (count && decoded > 0u + params.n_keep)
			llama_kv_cache_update(ctx);
		return true;
	};

	// Tokenize line and add to the tokens
	static auto push_str = [&params](const std::string& text, bool front = false, bool spec = false) -> uint {
		auto tt = llama_tokenize(model, text, false, spec);
		if (params.verbosity && !tt.empty())
			std::fprintf(stderr, "%s[tokens:%zu,%d] push %zu tokens: '%s'\n", g_esc.reset, tokens.size(), +front, tt.size(), text.c_str());
		if (front) {
			tokens.insert(tokens.begin() + params.n_keep + chunks.front(), tt.begin(), tt.end());
			chunks.front() += tt.size();
		} else {
			tokens.insert(tokens.end(), tt.begin(), tt.end());
			chunks.back() += tt.size();
		}
		return tt.size();
	};

	// Tokenize tr_text from id
	static auto push_id = [&params](line_id id, bool front = false) {
		auto tr_text = std::string_view(g_lines[id].tr_text);
		auto spos = tr_text.find("\n" + isuffix) + 1;
		if (!spos) {
			throw std::runtime_error("Line untranslated: " + g_lines[id].tr_text);
		}
		// Tokenize original strings with annotations
		std::size_t pref_pos = 0, post_pos = 0;
		if (!tr_text.starts_with(iprefix)) {
			post_pos = pref_pos = tr_text.find("\n" + iprefix) + 1;
			if (pref_pos == 0)
				throw std::runtime_error("Line without iprefix: " + g_lines[id].tr_text);
		}
		// Replay print_line logic
		auto out = apply_replaces(g_lines[id].text, false, 0);
		// Pre-annotations: no replaces, then original line with replaces
		g_lines[id].pre_ann = tr_text.substr(0, pref_pos);
		out = g_lines[id].pre_ann + iprefix + g_lines[id].name + std::move(out) + "\n";
		// Post-annotations: no replaces
		post_pos += iprefix.size();
		post_pos += g_lines[id].name.size();
		post_pos += g_lines[id].text.size() + 1;
		g_lines[id].post_ann = tr_text.substr(post_pos, spos - post_pos);
		out += g_lines[id].post_ann;
		out += isuffix;
		push_str(out, front);
		tr_text.remove_prefix(spos);
		tr_text.remove_prefix(isuffix.size());
		int token_count = 0;
		if (!g_lines[id].name.empty()) {
			// Tokenize name separately: find ": " delimiter
			auto pos = tr_text.find(": ") + 1;
			if (!pos) {
				throw std::runtime_error("Name untranslated: " + g_lines[id].tr_text);
			}
			if (isuffix.ends_with(" "))
				pos++;
			token_count += push_str(std::string(tr_text.substr(0, pos)), front);
			tr_text.remove_prefix(pos);
		}
		token_count += push_str(std::string(tr_text), front);
		g_lines[id].tr_tts.clear();
		g_lines[id].tr_tts.emplace_back(tokens.end() - token_count, tokens.end());
		if (token_count > params.n_predict + 1)
			throw std::runtime_error("Line too long: " + g_lines[id].tr_text);
		g_lines[id].tokens = front ? chunks.front() : chunks.back();
	};

	static auto load_cache = [&params](uint pos, uint hpos, line_id id) -> uint {
		auto fname = get_cache_file(s_cache_path, hpos, id);
		if (!fs::is_regular_file(fname))
			return 0;
		auto dummy = std::vector<int>(params.n_ctx);
		std::size_t count = -1;
		//auto s0 = std::chrono::steady_clock::now();
		if (!llama_state_seq_load_file(ctx, fname.c_str(), 1, dummy.data(), dummy.size(), &count)) {
			std::cerr << "Tokens: " << tokens.size() << std::endl;
			std::cerr << "Token count: " << llama_get_kv_cache_token_count(ctx) << std::endl;
			std::cerr << "Cells used: " << llama_get_kv_cache_used_cells(ctx) << std::endl;
			throw std::runtime_error("Failed to load KV cache: " + fname.string());
		}
		//std::cerr << (std::chrono::steady_clock::now() - s0).count() / 1e9 << "s" << std::endl;
		dummy.resize(count);
		if (auto expected = g_lines[id].tokens; expected != count)
			throw std::runtime_error("Bad KV cache size: " + fname.string());
		if (!std::equal(dummy.begin(), dummy.end(), tokens.begin() + pos)) {
			std::cerr << std::endl << "First mismatch pos: ";
			auto [mis1, mis2] = std::mismatch(dummy.begin(), dummy.end(), tokens.begin() + pos);
			std::cerr << (mis1 - dummy.begin()) << std::endl;
			std::cerr << "Expected tokens: '";
			for (auto it = mis2; it != tokens.begin() + pos + count; it++) {
				std::cerr << "[" << *it << "]";
				auto str = llama_token_to_piece(ctx, *it, false);
				REPLACE(str, "\n", "\\n");
				std::cerr << str;
			}
			std::cerr << "'" << std::endl << "  Loaded tokens: '";
			for (auto it = mis1; it != dummy.end(); it++) {
				std::cerr << "[" << *it << "]";
				auto str = llama_token_to_piece(ctx, *it, false);
				REPLACE(str, "\n", "\\n");
				std::cerr << str;
			}
			std::cerr << "'" << std::endl;
			throw std::runtime_error("Bad KV cache data: " + fname.string());
		}
		llama_kv_cache_seq_add(ctx, 1, 0, -1, +pos);
		llama_kv_cache_seq_add(ctx, 0, pos, -1, +count);
		llama_kv_cache_seq_cp(ctx, 1, 0, pos, pos + count);
		llama_kv_cache_seq_keep(ctx, 0);
		return count;
	};

	static auto decode_internal = [&params](uint count) -> void {
		if (params.verbosity) {
			std::cerr << "Decoding:" << g_esc.buf;
			auto end_it = tokens.begin() + decoded;
			if (tokens.end() - end_it > 300)
				end_it += 300;
			else
				end_it = tokens.end();
			for (auto it = tokens.begin() + decoded; it != end_it; it++) {
				auto str = llama_token_to_piece(ctx, *it, false);
				REPLACE(str, "\n", "\\n");
				std::cerr << str;
			}
			std::cerr << g_esc.reset << std::endl;
		}
		auto stamp0 = std::chrono::steady_clock::now();
		uint total = 0;
		if (count >= 0u + params.n_ctx)
			throw std::runtime_error("decode(): too many tokens: " + std::to_string(count));
		while (uint bsize = std::min<uint>(count, params.n_batch)) {
			// TODO: cannot properly interrupt by is_stopped, but probably not relevant anymore
			llama_decode(ctx, llama_batch_get_one(&tokens[decoded], bsize, decoded, 0));
			decoded += bsize;
			total += bsize;
			count -= bsize;
		}
		llama_synchronize(ctx);
		g_stats->raw_decodes += total;
		auto stamp1 = std::chrono::steady_clock::now();
		if (total > 1) {
			g_stats->batch_count += total;
			g_stats->batch_time += (stamp1 - stamp0).count() / 1000;
		}
	};

	static auto decode = [&params](uint injected = 0) -> void {
		if (tokens.size() == decoded)
			return;
		if (decoded == 0) {
			// Decode prompt first
			decode_internal(params.n_keep);
		}
		uint skip = params.n_keep;
		uint hpos = hist_pos;
		for (uint i = 0; i < injected; i++) {
			// Skip injected tokens at the beginning
			skip += chunks.at(i);
			decoded += chunks[i];
		}
		while (decoded < tokens.size()) {
			// Check if already decoded
			if (skip < decoded) {
				skip += chunks.at(injected + hpos - hist_pos);
				hpos++;
				continue;
			}
			// Load cached ids first
			if (std::lock_guard lock(g_mutex); hpos < g_history.size()) {
				if (auto res = load_cache(decoded, hpos, g_history[hpos])) {
					decoded += res;
					hpos++;
					skip = decoded;
					continue;
				}
			}
			// TODO: optimize by combining chunks
			decode_internal(chunks.at(injected + hpos - hist_pos));
			hpos++;
			skip = decoded;
		}
		if (llama_get_kv_cache_used_cells(ctx) + 0u != decoded) {
			throw std::runtime_error("used cells after decode");
		}
	};

	static auto inject_recollections = [&params](line_id id) -> uint {
		std::lock_guard lock(g_mutex);
		uint injected = 0;
		if (auto inj_list = get_inf_cache(params, id, hist_pos - 1); !inj_list.empty()) {
			for (auto it = inj_list.rbegin(); it != inj_list.rend(); it++) {
				// Push front
				chunks.emplace_front();
				push_id(g_history[*it], true);
				if (!load_cache(params.n_keep, *it, g_history[*it]))
					throw std::runtime_error("Unexpected: Cache not found");
				injected++;
			}
		}
		return injected;
	};

	static auto eject_recollections = [&params](uint injected) -> void {
		if (const uint inj_count = std::accumulate(chunks.begin(), chunks.begin() + injected, 0u)) {
			tokens.erase(tokens.begin() + params.n_keep, tokens.begin() + params.n_keep + inj_count);
			chunks.erase(chunks.begin(), chunks.begin() + injected);
			const auto p0 = params.n_keep;
			const auto p1 = params.n_keep + inj_count;
			if (!llama_kv_cache_seq_rm(ctx, 0, p0, p1))
				throw std::runtime_error("llama_kv_cache_seq_rm injected");
			llama_kv_cache_seq_add(ctx, 0, p1, -1, 0u - inj_count);
			decoded -= inj_count;
			if (llama_get_kv_cache_used_cells(ctx) + 0u != decoded)
				throw std::runtime_error("used cells injected");
		}
	};

	static auto init_segment = [&params]() -> void {
		// Initialize segment: eject all first
		eject_bunch(-1);

		// Load full history
		for (auto& id : g_history) {
			if (g_lines[id].tr_text.empty())
				break;
			chunks.emplace_back();
			push_id(id);
			// Check cache file
			auto fname = get_cache_file(s_cache_path, &id - g_history.data(), id);
			if (!fs::is_regular_file(fname)) {
				std::cerr << "Cache not found: " << fname << std::endl;
				if (!eject_start())
					throw std::runtime_error("eject_start init seg");
				decode();
				//const uint injected = inject_recollections(id);
				//decode(injected);
				//eject_recollections(injected);
			}
		}

		if (!eject_start())
			throw std::runtime_error("eject_start init_seg");
		decode();
	};

	auto add_tail_finalize = [id]() -> void {
		if (segment + 1) {
			auto& tail = g_lines.segs[segment].tr_tail;
			if (id.second == 0 && !g_lines.segs[segment].lines.back().tr_text.empty()) {
				// Finalize segment if necessary
				if (tail.empty()) {
					tail = "\n";
					update_segment(segment, true, -1);
				}
			}
		}
	};

	if (cmd == tr_cmd::eject) {
		eject_bunch(id.second);
		return true;
	}

	if (cmd == tr_cmd::unload) {
		unload();
		return true;
	}

	if (cmd == tr_cmd::reload) {
		if (id == c_bad_id) {
			// Full reload
			init_segment();
			return true;
		}
		if (id.first != segment) {
			add_tail_finalize();
			segment = id.first;
		}
		if (segment + 1) {
			for (line_id nid{segment, id.second}; g_lines.is_valid(nid); g_lines.advance(nid)) {
				if (g_lines[nid].tr_text.empty())
					break;
				chunks.emplace_back();
				push_id(nid);
				// TODO: allow restoring linear history as well, not just ejecting it
				// Otherwise, reloaded lines are a bit inconsistent with the rest.
				if (!eject_start())
					throw std::runtime_error("eject_start init seg");
				//const uint injected = inject_recollections(id);
				//decode(injected);
				//eject_recollections(injected);
				decode();
			}
		}
		return true;
	}

	if (id.first != segment) {
		// Update previous translation file
		if (segment + 1) {
			add_tail_finalize();
		}

		segment = id.first;
		if (tokens.empty())
			init_segment();
	}

	// Translate line(s)
	if (id != c_bad_id) {
		// Discard current and following lines
		uint to_eject = 0;
		for (line_id nid = id; cmd == tr_cmd::translate && nid != c_bad_id; g_lines.advance(nid)) {
			std::lock_guard lock(g_mutex);

			if (g_lines[nid].tr_text.empty())
				break;
			to_eject++;
			if (params.verbosity)
				std::fprintf(stderr, "%s[id:%u:%u] Ejected\n", g_esc.reset, nid.first, nid.second);
			g_lines[nid].tr_text = {};
			if (g_lines[nid].tr_tts.empty())
				throw std::runtime_error("tr_tts not found: " + g_lines[nid].text);
			// tr_tts is not discarded here
		}
		// TODO: more accurate computation with accounting for context size
		if (to_eject >= 10) {
			// Full reset
			init_segment();
		} else {
			eject_bunch(to_eject);
		}
	} else {
		return true;
	}

	// Detect additional lines for translation
	bool echo = false;
	for (line_id pid{segment, 0}; cmd == tr_cmd::translate && pid <= id; g_lines.advance(pid)) {
		auto& line = g_lines[pid];
		if (!(std::lock_guard{g_mutex}, line.tr_text.empty()))
			continue;

		std::string llama_out;
		llama_out += line.pre_ann;
		// TODO: there is a race condition here when tr_tts is accidentally not empty (hopefully not a data race)
		if (pid < id || line.tr_tts.empty())
			echo = true; // Sticky flag to display source line/annotations, otherwise this is rewrite request
		if (std::this_thread::get_id() == s_main_tid && echo)
			std::cout << g_esc.orig << line.pre_ann << std::flush;
		std::string spker = print_line(pid, &llama_out, std::this_thread::get_id() == s_main_tid && echo);
		if (std::this_thread::get_id() == s_main_tid && echo)
			std::cout << g_esc.orig << line.post_ann << std::flush;
		llama_out += line.post_ann;
		llama_out += isuffix;
		chunks.emplace_back();
		push_str(llama_out);
		// Encode speaker separately and count its tokens
		int pred_count = 0;
		// Process rewrite request (cfg_negative_prompt is just abused to pass the selection)
		if (std::this_thread::get_id() == s_main_tid && !g_neg.empty()) {
			// Copy previous translation's tokens preceding selection
			llama_out.clear();
			std::vector<std::string> tstrs;
			for (auto t : line.tr_tts.back()) {
				tstrs.emplace_back(llama_token_to_piece(ctx, t));
				llama_out += tstrs.back();
				pred_count++;
				if (auto negpos = llama_out.find(g_neg); negpos + 1) {
					while (llama_out.size() > negpos) {
						llama_out -= tstrs.back();
						tstrs.pop_back();
						pred_count--;
					}
					while (!isuffix.ends_with(" ") && llama_out.ends_with(" ")) {
						llama_out -= tstrs.back();
						tstrs.pop_back();
						pred_count--;
					}
					break;
				}
			}
			if (params.verbosity)
				std::fprintf(stderr, "%s[tokens:%zu] push %zu tokens (repeated)\n", g_esc.reset, tokens.size(), tstrs.size());
			tokens.insert(tokens.end(), line.tr_tts.back().begin(), line.tr_tts.back().begin() + tstrs.size());
			chunks.back() += pred_count;
		} else {
			pred_count += push_str(spker);
			llama_out = spker;
		}
		if (g_mode == op_mode::rt_llama && std::this_thread::get_id() == s_main_tid) {
			std::cout << g_esc.tran << llama_out.c_str() + llama_out.starts_with(" ");
			std::cout << std::flush; // Stream to terminal
		}
		if (!eject_start())
			return false;

		if (params.verbosity)
			std::fprintf(stderr, "%s[id:%u:%u, t:%zu, last:%u, chk:%llu] Line: %s%s\n", g_esc.reset, pid.first, pid.second, tokens.size(),
						 chunks.size() ? chunks.back() : -1, std::accumulate(tokens.begin(), tokens.end(), 0ull), line.name.c_str(),
						 line.text.data());

		// Inject context temporarily, decode pending tokens
		const uint injected = inject_recollections(pid);
		decode(injected);

		auto sparams = params.sparams;
		if (!line.tr_tts.empty())
			sparams.seed += std::accumulate(line.tr_tts.back().begin(), line.tr_tts.back().end(), line.tr_tts.size());
		const auto sctx = gpt_sampler_init(model, sparams);
		const bool use_grammar = !line.name.empty();
		for (int i = 0; i < pred_count; i++)
			gpt_sampler_accept(sctx, *(tokens.rbegin() + pred_count - i), use_grammar);
		static const auto real_isuffix = "\n" + isuffix;
		const auto replaces = make_name_replaces(line.text);
		std::size_t last_neg = spker.size();
		std::size_t last_suf = llama_out.size();
		while (!is_stopped() && !llama_out.ends_with("\n")) {
			// Predict next token
			auto stamp0 = std::chrono::steady_clock::now();
			auto logits = llama_get_logits(ctx);
			for (auto& vec : line.tr_tts) {
				if (pred_count + 0u >= vec.size())
					continue;
				const auto token_str = llama_token_to_piece(ctx, vec[pred_count]);

				// Workaround for not penalizing certain patterns
				static constexpr std::string_view s_no_penalty[]{
					" \"",
					" '",
					" (",
					")",
				};
				if ((!line.name.empty() && llama_out == spker) || pred_count == 0) {
					bool found = false;
					for (auto s : s_no_penalty) {
						if (llama_out != spker) {
							while (s.starts_with(" "))
								s.remove_prefix(1);
						}
						if (token_str.starts_with(s)) {
							found = true;
							break;
						}
					}
					if (found)
						break;
				}
				if (std::equal(tokens.end() - pred_count, tokens.end(), vec.begin())) {
					// Don't penalize pure space if it appeared here before
					auto tstr = llama_token_to_piece(ctx, vec[pred_count], false);
					if (!isuffix.ends_with(" ") && token_str == " ")
						continue;
					// Penalize specific token (except newline)
					if (tstr != "\n")
						logits[vec[pred_count]] -= 1.f;
					// Penalize pure space (experimentally observed to appear sometimes and break the line)
					static const auto s_tspace = llama_tokenize(ctx, " ", false);
					logits[s_tspace[0]] -= 1.f;
				}
			}
			auto token_id = gpt_sampler_sample(sctx, ctx, -1);
			g_stats->raw_samples++;
			auto stamp1 = std::chrono::steady_clock::now();
			g_stats->sample_time += (stamp1 - stamp0).count() / 1000;
			static const auto s_tnl = llama_tokenize(ctx, "\n", false);
			if (++pred_count > params.n_predict)
				token_id = s_tnl[0]; // Force newline if size exceeded
			if (llama_token_is_eog(model, token_id))
				token_id = s_tnl[0]; // Force newline on EOT/EOS
			auto token_str = llama_token_to_piece(ctx, token_id);
			if (token_str == "\n")
				token_id = s_tnl[0]; // Fix alternative newline (TODO: correct workaround)
			if (std::count(token_str.begin(), token_str.end(), '\n') > 0u + token_str.ends_with("\n")) {
				// Attempt to fix the line containing incorrect newlines
				token_str.resize(token_str.find_first_of('\n') + 1);
				auto tks = llama_tokenize(model, token_str, false);
				if (tks.size() == 1)
					token_id = tks[0];
				else {
					std::cerr << "Token sequence was forsibly converted to newline: " << token_str;
					token_id = s_tnl[0];
				}
			}
			tokens.push_back(token_id);

			gpt_sampler_accept(sctx, token_id, use_grammar);
			llama_out += token_str;
			int to_decode = 1;
			int to_penalize = -1;
			auto print_tokens = [&]() {
				if (!params.verbosity)
					return;
				std::fprintf(stderr, "%sT: ", g_esc.reset);
				for (auto it = tokens.end() - pred_count; it != tokens.end(); it++) {
					std::fprintf(stderr, "%s|", llama_token_get_text(model, *it));
				}
				std::fprintf(stderr, " (dec:%d,pen='%s')\n", to_decode, to_penalize >= 0 ? llama_token_get_text(model, to_penalize) : "");
			};

			// Try to parse translated name
			if (spker.empty() && !line.name.empty()) {
				std::string_view speaker = llama_out;
				if (auto pos = speaker.find_first_not_of(" ") + 1)
					speaker.remove_prefix(pos - 1);
				speaker = speaker.substr(0, speaker.find(":") + 1);
				if (!speaker.empty()) {
					std::lock_guard lock(g_mutex);
					spker = g_speakers[line.name] = speaker;
				}
			}
			if (!spker.empty() || line.name.empty()) {
				// TODO: implement separate replace system whith takes original text into account
				auto fixed = llama_out;
				// Also apply "neg" as a replacement
				if (auto nf = fixed.find(g_neg, last_neg); !g_neg.empty() && nf + 1) {
					// Completely cut the continuation
					fixed.resize(nf);
					// Sanitize trailing spaces because it can break the output
					while (!isuffix.ends_with(" ") && fixed.ends_with(" ")) {
						fixed -= " ";
					}
					last_neg = nf + 1;
				}
				// Apply name suffix replaces as well
				for (auto& [from, to] : replaces) {
					if (from == to) [[unlikely]] {
						// Special case (name was found without any known suffix)
						if (llama_out.ends_with(from) && llama_out == fixed) {
							// Penalize tokens starting with "-"
							to_penalize = -2;
							break;
						}
					} else if (auto nf = fixed.find(from, last_suf); nf + 1) {
						// Completely cut the continuation
						last_suf = nf + to.size();
						fixed.resize(nf);
						fixed += to;
						break;
					}
				}
				// Retokenize to also make the tokenization uniform
				auto tt = llama_tokenize(model, fixed, false);
				if (!fixed.empty() && (fixed != llama_out || !std::equal(tt.begin(), tt.end(), tokens.end() - pred_count))) {
					if (g_mode == op_mode::rt_llama && std::this_thread::get_id() == s_main_tid) {
						// Overwrite current line
						if (fixed.size() < llama_out.size()) {
							std::cout << std::string(llama_out.size() - fixed.size(), ' ');
							std::cout << std::string(llama_out.size() - fixed.size(), '\b');
						}
						std::cout << std::string(llama_out.size(), '\b') << g_esc.tran << fixed.c_str() + (fixed.starts_with(" "));
						std::cout << std::flush;
					}
					// Replace only the few last tokens
					auto [mis, mistt] = std::mismatch(tokens.end() - pred_count, tokens.end(), tt.begin(), tt.end());
					uint miscount = tokens.end() - mis;
					int dropped = mis != tokens.end();
					if (mistt == tt.end()) {
						// If nothing to replace, re-evaluate last token and penalize token that was following it
						// TODO: better logic
						if (llama_token_to_piece(ctx, *mis, false) != " ")
							to_penalize = *mis;
						to_decode++;
					}
					tokens.erase(mis, tokens.end());
					int to_remove = miscount - dropped;
					if (mistt == tt.end()) {
						to_remove++;
					}
					to_decode += tt.end() - mistt;
					to_decode -= dropped;
					tokens.insert(tokens.end(), mistt, tt.end());
					chunks.back() -= to_remove;
					const auto p0 = decoded - to_remove;
					decoded -= to_remove;
					if (!llama_kv_cache_seq_rm(ctx, 0, p0, -1))
						throw std::runtime_error("llama_kv_cache_seq_rm replaces");
					gpt_sampler_reset(sctx);
					for (auto& t : tt)
						gpt_sampler_accept(sctx, t, use_grammar);
					llama_out = std::move(fixed);
					pred_count += to_decode - to_remove - 1;
					token_str.clear();
				}
			}
			print_tokens();
			if (g_mode == op_mode::rt_llama && std::this_thread::get_id() == s_main_tid && !token_str.empty()) {
				std::cout << token_str.c_str() + (token_str.starts_with(" ") && llama_out == token_str);
				std::cout << std::flush; // Stream to terminal
			}

			if (tokens.size() >= 0u + params.n_ctx)
				throw std::runtime_error("eval: too many tokens: " + std::to_string(tokens.size()));
			if (decoded >= tokens.size())
				throw std::runtime_error("eval: bad pointer: " + std::to_string(decoded));
			stamp0 = std::chrono::steady_clock::now();
			if (auto result = llama_decode(ctx, llama_batch_get_one(&tokens[decoded], to_decode, decoded, 0))) {
				std::cerr << "llama_decode failed: " << result << std::endl;
				return false;
			}
			llama_synchronize(ctx);
			g_stats->raw_decodes += to_decode;
			stamp1 = std::chrono::steady_clock::now();
			g_stats->eval_time += (stamp1 - stamp0).count() / 1000;
			decoded += to_decode;
			if (chunks.empty())
				throw std::out_of_range("chunks.empty()");
			chunks.back() += to_decode;
			if (decoded > 0u + params.n_ctx)
				throw std::out_of_range("ctx overflow");
			logits = llama_get_logits(ctx);
			if (to_penalize >= 0) {
				logits[to_penalize] -= 1.f;
			} else if (to_penalize == -2) {
				for (int i = 0; i < llama_n_vocab(model); i++) {
					if (llama_token_get_text(model, i)[0] == '-') {
						logits[i] -= 1.f;
					}
				}
			}
		}

		// Remove injected tokens
		eject_recollections(injected);
		gpt_sampler_free(sctx);
		g_stats->sample_count += pred_count;
		if (is_stopped() && !llama_out.ends_with("\n") && std::this_thread::get_id() == s_main_tid)
			std::cout << std::endl; // Stop current line
		if (g_mode == op_mode::rt_llama && std::this_thread::get_id() == s_main_tid)
			std::cout << g_esc.reset << std::flush; // Reset terminal colors
		if (is_stopped() && !llama_out.ends_with("\n")) {
			// Discard incomplete work
			eject_bunch(1, false);
			return true;
		}

		if (spker.empty() && !line.name.empty())
			std::cerr << "Failed to parse speaker translation for: " << line.name << std::endl;

		// Store translation tokens for possible reuse
		line.tr_tts.emplace_back(tokens.end() - pred_count, tokens.end());

		// Store translated line
		std::lock_guard lock(g_mutex);
		line.tr_text.clear();
		line.tr_text += line.pre_ann;
		line.tr_text += iprefix;
		line.tr_text += line.name;
		line.tr_text += line.text;
		line.tr_text += '\n';
		line.tr_text += line.post_ann;
		line.tr_text += isuffix;
		line.tr_text += llama_out;
		line.tokens = chunks.back();
		g_stats->raw_translates++;
	}

	if (!is_stopped() && std::this_thread::get_id() == s_main_tid && g_mode == op_mode::rt_llama && id != c_bad_id) {
		// Launch worker thread to continue translation in the background
		// Check if thread is disabled:
		if (params.n_draft <= 0)
			return true;
		// Check if the next line is untranslated
		if (auto next = g_lines.next(id); next != c_bad_id) {
			std::lock_guard lock(g_mutex);
			if (!g_lines[next].tr_text.empty())
				return true;
		} else {
			// Nothing to do
			return true;
		}

		stop_sig = -1;
		work_res = id.second + 1;
		worker = std::thread([=, &params] {
			std::deque<uint> old_pos;
			old_pos.push_back(hist_pos);
			auto start = g_lines.next(id);
			auto next = start;
			auto last = id;
			for (int i = 0; i < params.n_draft; i++) {
				last.second++;
				if (g_lines.is_last(last))
					break;
			}
			uint ahead = 0;
			if (params.verbosity)
				std::fprintf(stderr, "%s[id:%u:%u] Thread entry\n", g_esc.reset, next.first, next.second);
			while (!is_stopped() && next != c_bad_id) {
				if (next <= last) {
					if (!translate(params, next, tr_cmd::translate)) {
						// No space left in context
						throw std::runtime_error("Unexpected: no space in context");
					}
					work_res++;
					g_lines.advance(next);
					ahead += 1;
					old_pos.push_back(hist_pos);
					continue;
				}
				std::unique_lock lock(g_mutex);
				while (!is_stopped() && next > last) {
					ahead = last.second - g_history.back().second;
					// Fix negative values, should mean untranslated lines were added to history
					for (auto it = g_history.rbegin(); it != g_history.rend(); it++) {
						if (!g_lines[*it].tr_text.empty())
							break;
						ahead++;
					}
					if (ahead > INT_MAX)
						throw std::runtime_error("Unexpected ahead: " + std::to_string(ahead));
					for (int i = 0; i < params.n_draft; i++) {
						if (ahead + i < params.n_draft + 0u) {
							last.second++;
							if (g_lines.is_last(last))
								break;
						} else
							break;
					}
					if (next <= last)
						break;
					work_cv.wait(lock);
				}
			}
			// Wait for thread's destiny instead of terminating it immediately
			std::unique_lock lock(g_mutex);
			while (stop_sig == 0u - 1)
				work_cv.wait(lock);
			// Discard translations which were done in the background
			bool msg = false;
			uint from = stop_sig.load();
			while (start != c_bad_id && !g_lines[start].tr_text.empty()) {
				if (start.second >= from) {
					if (params.verbosity && !std::exchange(msg, true))
						std::fprintf(stderr, "%s[id:%u:%u, from:%u] Cleaning\n", g_esc.reset, start.first, start.second, from);
					g_lines[start].tr_text.clear();
					g_lines[start].tr_tts.clear();
					eject_bunch(1); // reverse order but should work ok?
					old_pos.pop_back();
				}
				g_lines.advance(start);
			}
			// Restore loaded history
			if (llama_get_kv_cache_used_cells(ctx) + 0u != decoded) {
				throw std::runtime_error("used cells before rollback");
			}
			while (hist_pos > old_pos.back()) {
				auto hid = g_history[--hist_pos];
				chunks.emplace_front();
				push_id(hid, true);
				auto res = load_cache(params.n_keep, hist_pos, hid);
				if (!res)
					throw std::runtime_error("Unexpected: cache not found");
				decoded += res;
				if (llama_get_kv_cache_used_cells(ctx) + 0u != decoded) {
					throw std::runtime_error("used cells after rollback");
				}
			}
			if (params.verbosity)
				std::fprintf(stderr, "%s[id:%u:?] Thread exit\n", g_esc.reset, next.first);
		});
	}

	return true;
}