Files
QtPotato/src/translate/translationservice.cpp
T

626 lines
20 KiB
C++

#include "translationservice.h"
#include "apikeystore.h"
#include "qtpotatosettings.h"
#include "translationpresets.h"
#include <KLocalizedString>
#include <QCoroFuture>
#include <QCoroTimer>
#include <QCryptographicHash>
#include <QDir>
#include <QFileInfo>
#include <QHash>
#include <QLocale>
#include <QNetworkAccessManager>
#include <QPointer>
#include <QSaveFile>
#include <QStandardPaths>
#include <QtConcurrentRun>
#include <chrono>
#include <vector>
using namespace Qt::Literals::StringLiterals;
using namespace std::chrono_literals;
namespace
{
// Source lines before a batch that the model sees for context.
constexpr qsizetype ContextLines = 6;
constexpr int MaxAttempts = 3;
// Requests in flight at once. Batches only use source lines as context, so they don't
// depend on each other.
constexpr int ParallelRequests = 3;
// Bump when the prompt changes so old cached translations are not reused.
constexpr auto CacheVersion = "llm-v2";
// "Vietnamese" -> "vi", so the file ends in ".vi.srt" like other players expect.
QString languageSuffix(const QString &language)
{
for (int i = QLocale::C + 1; i <= QLocale::LastLanguage; ++i) {
const auto candidate = static_cast<QLocale::Language>(i);
if (QLocale::languageToString(candidate).compare(language, Qt::CaseInsensitive) == 0) {
const QString code = QLocale::languageToCode(candidate, QLocale::ISO639Part1);
return code.isEmpty() ? QLocale::languageToCode(candidate) : code;
}
}
QString suffix = language.toLower();
suffix.removeIf([](QChar c) {
return !c.isLetterOrNumber() && c != u'-';
});
return suffix.isEmpty() ? u"translated"_s : suffix;
}
QString cachePath(const QUrl &videoUrl, const QString &key, const QString &language)
{
const QString dir = QStandardPaths::writableLocation(QStandardPaths::CacheLocation) + u"/translations"_s;
const QString baseName = videoUrl.isLocalFile() ? QFileInfo(videoUrl.toLocalFile()).completeBaseName() : u"subtitles"_s;
return u"%1/%2.%3.%4.srt"_s.arg(dir, baseName, key, languageSuffix(language));
}
}
struct TranslationService::Job {
std::atomic_bool cancelled = false;
std::unique_ptr<LlmClient> client;
QString mediaTitle;
};
struct TranslationService::Batches {
QStringList lines;
QStringList translations;
qsizetype batchSize = 0;
qsizetype nextStart = 0;
qsizetype doneLines = 0;
QString error;
};
TranslationService::TranslationService(QObject *parent)
: QObject(parent)
, m_network(new QNetworkAccessManager(this))
{
}
TranslationService::~TranslationService()
{
cancel();
}
bool TranslationService::isBusy() const
{
return m_busy;
}
double TranslationService::progress() const
{
return m_progress;
}
QString TranslationService::statusText() const
{
return m_statusText;
}
QString TranslationService::errorText() const
{
return m_errorText;
}
QStringList TranslationService::availableModels() const
{
return m_availableModels;
}
QVariantList TranslationService::translationPresets() const
{
QVariantList result;
for (const auto &preset : TranslationPresets::all()) {
result.append(QVariantMap{{u"id"_s, preset.id}, {u"name"_s, preset.name}, {u"description"_s, preset.description}, {u"guidance"_s, preset.guidance}});
}
return result;
}
bool TranslationService::isFetchingModels() const
{
return m_fetchingModels;
}
QVariantList TranslationService::glossaryEntries() const
{
QVariantList result;
for (const auto &entry : m_glossary.entries) result.append(QVariantMap{{u"source"_s, entry.source}, {u"target"_s, entry.target}});
return result;
}
QString TranslationService::glossaryDirectory() const
{
return m_glossary.directory;
}
QString TranslationService::glossaryError() const
{
return m_glossaryError;
}
void TranslationService::loadGlossary(const QUrl &videoUrl)
{
const QString directory = videoUrl.isLocalFile() ? QFileInfo(videoUrl.toLocalFile()).absolutePath() : QString();
m_glossary = directory.isEmpty() ? GlossaryStore::Glossary{} : GlossaryStore::load(directory);
if (!m_glossaryError.isEmpty()) {
m_glossaryError.clear();
Q_EMIT glossaryErrorChanged();
}
Q_EMIT glossaryChanged();
}
bool TranslationService::saveGlossary(const QVariantList &entries)
{
QList<GlossaryStore::Entry> result;
for (const QVariant &value : entries) {
const QVariantMap map = value.toMap();
result.append({map.value(u"source"_s).toString().trimmed(), map.value(u"target"_s).toString().trimmed()});
}
GlossaryStore::Glossary candidate = m_glossary;
candidate.entries = result;
if (!candidate.directory.isEmpty()) {
QString error;
if (!GlossaryStore::save(candidate, &error)) {
const QString message = error.isEmpty() ? i18n("Could not save the translation glossary.")
: i18n("Could not save the translation glossary: %1", error);
if (m_glossaryError != message) {
m_glossaryError = message;
Q_EMIT glossaryErrorChanged();
}
return false;
}
}
m_glossary = candidate;
if (!m_glossaryError.isEmpty()) {
m_glossaryError.clear();
Q_EMIT glossaryErrorChanged();
}
Q_EMIT glossaryChanged();
return true;
}
void TranslationService::clearGlossary()
{
m_glossary.entries.clear();
if (!m_glossary.directory.isEmpty()) {
GlossaryStore::save(m_glossary);
}
Q_EMIT glossaryChanged();
}
void TranslationService::translate(const QUrl &videoUrl, const QVariantMap &track, const QString &mediaTitle)
{
cancel();
m_job = std::make_shared<Job>();
m_job->mediaTitle = mediaTitle;
run(videoUrl, track);
}
void TranslationService::cancel()
{
if (!m_job) {
return;
}
// The running coroutine keeps its own reference and stops at its next check.
m_job->cancelled = true;
if (m_job->client) {
m_job->client->abort();
}
m_job.reset();
setBusy(false);
setStatusText({});
}
void TranslationService::fetchModels()
{
if (!m_fetchingModels) {
runFetchModels();
}
}
void TranslationService::clearError()
{
setErrorText({});
}
QCoro::Task<std::pair<LlmClient::Config, QString>> TranslationService::loadConfig()
{
const auto settings = QtPotatoSettings::self();
LlmClient::Config config{
LlmClient::normalizedBaseUrl(settings->apiBaseUrl()),
{},
settings->model().trimmed(),
settings->targetLanguage().trimmed(),
settings->translationPreset().trimmed(),
settings->customInstructions(),
{},
};
if (!config.baseUrl.isValid() || config.baseUrl.scheme().isEmpty() || config.baseUrl.host().isEmpty()) {
co_return std::pair{config, i18n("Set the API base URL in Settings first.")};
}
const ApiKeyStore::ReadResult key = co_await ApiKeyStore::readKey();
if (!key.error.isEmpty()) {
co_return std::pair{config, i18n("Could not read the API key from the wallet: %1", key.error)};
}
config.apiKey = key.key;
co_return std::pair{config, QString()};
}
QCoro::Task<Subtitles::ExtractResult> TranslationService::loadCues(std::shared_ptr<Job> job, QUrl videoUrl, QVariantMap track)
{
if (track.value(u"external"_s).toBool()) {
const QString path = track.value(u"externalFilename"_s).toString();
co_return co_await QtConcurrent::run([path] {
QString error;
QList<SubtitleCue> cues = Subtitles::parseFile(path, &error);
return Subtitles::ExtractResult{cues, error};
});
}
if (!videoUrl.isLocalFile()) {
co_return Subtitles::ExtractResult{{}, i18n("Only subtitles of local files can be translated.")};
}
const int streamIndex = track.value(u"ffIndex"_s, -1).toInt();
if (streamIndex < 0) {
co_return Subtitles::ExtractResult{{}, i18n("mpv did not report where this subtitle track is stored.")};
}
const QString path = videoUrl.toLocalFile();
co_return co_await QtConcurrent::run([path, streamIndex, job] {
return Subtitles::extractEmbedded(path, streamIndex, job->cancelled);
});
}
QCoro::Task<> TranslationService::run(QUrl videoUrl, QVariantMap track)
{
const std::shared_ptr<Job> job = m_job;
setErrorText({});
setProgress(0);
setBusy(true);
setStatusText(i18n("Reading subtitles…"));
auto [config, configError] = co_await loadConfig();
if (job->cancelled) {
co_return;
}
if (!configError.isEmpty()) {
fail(job, configError);
co_return;
}
if (config.model.isEmpty()) {
fail(job, i18n("Choose a model in Settings first."));
co_return;
}
Subtitles::ExtractResult source = co_await loadCues(job, videoUrl, track);
if (job->cancelled) {
co_return;
}
if (!source.error.isEmpty()) {
fail(job, source.error);
co_return;
}
const QList<SubtitleCue> cues = Subtitles::mergeRepeated(source.cues);
if (cues.isEmpty()) {
fail(job, i18n("This subtitle track has no text to translate."));
co_return;
}
// Episodes in one folder share a glossary. Streams have no folder, so they get none.
const QString seriesDirectory = videoUrl.isLocalFile() ? QFileInfo(videoUrl.toLocalFile()).absolutePath() : QString();
const bool useGlossary = !seriesDirectory.isEmpty();
GlossaryStore::Glossary glossary = useGlossary ? GlossaryStore::load(seriesDirectory) : GlossaryStore::Glossary{};
const QString sourceSrt = Subtitles::toSrt(cues);
// Same subtitles, model, language, prompt settings and glossary give the same file, so it is translated once.
const auto outputPathFor = [&](const GlossaryStore::Glossary &g) {
QCryptographicHash hash(QCryptographicHash::Sha1);
hash.addData(CacheVersion);
hash.addData(config.model.toUtf8());
hash.addData(config.targetLanguage.toUtf8());
hash.addData(config.translationPreset.toUtf8());
hash.addData(config.customInstructions.toUtf8());
hash.addData(GlossaryStore::serialize(g).toUtf8());
hash.addData(sourceSrt.toUtf8());
return cachePath(videoUrl, QString::fromLatin1(hash.result().toHex().left(12)), config.targetLanguage);
};
QString outputPath = outputPathFor(glossary);
if (QFileInfo::exists(outputPath)) {
m_job.reset();
setBusy(false);
setStatusText({});
Q_EMIT translated(QUrl::fromLocalFile(outputPath), true);
co_return;
}
if (useGlossary && QtPotatoSettings::self()->autoBuildGlossary()) {
setStatusText(i18n("Building glossary…"));
QStringList extractionLines;
QHash<QString, bool> seen;
for (const SubtitleCue &cue : cues) {
if (!seen.contains(cue.text)) { seen.insert(cue.text, true); extractionLines.append(cue.text); }
}
constexpr qsizetype MaxGlossaryChars = 30000;
// Long subtitles are sampled evenly so names from the whole episode are seen.
qsizetype allChars = 0;
for (const QString &line : std::as_const(extractionLines)) {
allChars += line.size() + 1;
}
const qsizetype step = std::max<qsizetype>(1, (allChars + MaxGlossaryChars - 1) / MaxGlossaryChars);
QStringList sampled;
qsizetype total = 0;
for (qsizetype i = 0; i < extractionLines.size(); i += step) {
if (total + extractionLines.at(i).size() + 1 > MaxGlossaryChars) {
break;
}
sampled.append(extractionLines.at(i));
total += extractionLines.at(i).size() + 1;
}
job->client = std::make_unique<LlmClient>(m_network, config);
const LlmClient::BatchResult extracted = co_await job->client->extractGlossary(sampled, job->mediaTitle);
if (job->cancelled) {
co_return;
}
if (extracted.error.isEmpty()) {
QList<GlossaryStore::Entry> additions;
for (const QString &encoded : extracted.translations) {
const qsizetype tab = encoded.indexOf(u'\t');
if (tab > 0) additions.append({encoded.left(tab), encoded.mid(tab + 1)});
}
glossary = GlossaryStore::merge(glossary, additions);
GlossaryStore::save(glossary);
} else {
setStatusText(i18n("Glossary unavailable; translating without it."));
}
outputPath = outputPathFor(glossary);
if (QFileInfo::exists(outputPath)) {
m_job.reset();
setBusy(false);
setStatusText({});
Q_EMIT translated(QUrl::fromLocalFile(outputPath), true);
co_return;
}
}
// Repeated lines (songs, catchphrases) are sent once.
QStringList uniqueLines;
QHash<QString, qsizetype> uniqueIndex;
QList<qsizetype> cueToUnique;
for (const SubtitleCue &cue : cues) {
auto it = uniqueIndex.constFind(cue.text);
if (it == uniqueIndex.cend()) {
it = uniqueIndex.insert(cue.text, uniqueLines.size());
uniqueLines.append(cue.text);
}
cueToUnique.append(*it);
}
QStringList glossaryLines;
for (const auto &entry : std::as_const(glossary.entries)) {
glossaryLines.append(entry.source + u" → "_s + entry.target);
}
config.glossaryText = glossaryLines.join(u'\n');
job->client = std::make_unique<LlmClient>(m_network, config);
auto batches = std::make_shared<Batches>();
batches->lines = uniqueLines;
batches->translations.resize(uniqueLines.size());
batches->batchSize = QtPotatoSettings::self()->batchSize();
setStatusText(i18n("Translating… %1 of %2 lines", 0, uniqueLines.size()));
// Tasks start running as soon as they are created; awaiting them just waits for the end.
std::vector<QCoro::Task<>> workers;
for (int i = 0; i < ParallelRequests; ++i) {
workers.push_back(translateBatches(job, batches));
}
for (QCoro::Task<> &worker : workers) {
co_await std::move(worker);
}
if (job->cancelled) {
co_return;
}
if (!batches->error.isEmpty()) {
fail(job, batches->error);
co_return;
}
QList<SubtitleCue> output = cues;
for (qsizetype i = 0; i < output.size(); ++i) {
output[i].text = batches->translations.at(cueToUnique.at(i));
}
QDir().mkpath(QFileInfo(outputPath).absolutePath());
QSaveFile file(outputPath);
if (!file.open(QIODevice::WriteOnly) || file.write(Subtitles::toSrt(output).toUtf8()) < 0 || !file.commit()) {
fail(job, i18n("Could not save the translation to %1: %2", outputPath, file.errorString()));
co_return;
}
m_job.reset();
setBusy(false);
setStatusText({});
Q_EMIT translated(QUrl::fromLocalFile(outputPath), false);
}
QCoro::Task<LlmClient::BatchResult> TranslationService::translateLines(std::shared_ptr<Job> job, QStringList lines, QStringList context)
{
LlmClient::BatchResult result;
for (int attempt = 1; attempt <= MaxAttempts; ++attempt) {
result = co_await job->client->translate(lines, context, job->mediaTitle);
if (job->cancelled || result.error.isEmpty() || !result.retryable || attempt == MaxAttempts) {
break;
}
co_await QCoro::sleepFor(2s * attempt);
if (job->cancelled) {
co_return result;
}
}
if (job->cancelled || !result.error.isEmpty()) {
co_return result;
}
QList<qsizetype> missing;
for (qsizetype i = 0; i < result.translations.size(); ++i) {
if (result.translations.at(i).isNull()) {
missing.append(i);
}
}
if (missing.isEmpty()) {
co_return result;
}
if (lines.size() == 1) {
co_return LlmClient::BatchResult{{}, i18n("The model did not translate this line: %1", lines.first()), false};
}
if (missing.size() == lines.size()) {
// Nothing came back usable; smaller requests are easier for the model.
const qsizetype half = lines.size() / 2;
const LlmClient::BatchResult first = co_await translateLines(job, lines.first(half), context);
if (job->cancelled || !first.error.isEmpty()) {
co_return first;
}
QStringList secondContext = context + lines.first(half);
secondContext = secondContext.mid(std::max<qsizetype>(0, secondContext.size() - ContextLines));
const LlmClient::BatchResult second = co_await translateLines(job, lines.sliced(half), secondContext);
if (job->cancelled || !second.error.isEmpty()) {
co_return second;
}
co_return LlmClient::BatchResult{first.translations + second.translations, {}, false};
}
// Ask again only for the lines the model skipped.
QStringList retryLines;
for (const qsizetype i : missing) {
retryLines.append(lines.at(i));
}
const LlmClient::BatchResult retry = co_await translateLines(job, retryLines, context);
if (job->cancelled || !retry.error.isEmpty()) {
co_return retry;
}
for (qsizetype i = 0; i < missing.size(); ++i) {
result.translations[missing.at(i)] = retry.translations.at(i);
}
co_return result;
}
QCoro::Task<> TranslationService::translateBatches(std::shared_ptr<Job> job, std::shared_ptr<Batches> batches)
{
// After an error, requests already sent may finish, but no new ones start.
while (!job->cancelled && batches->error.isEmpty() && batches->nextStart < batches->lines.size()) {
const qsizetype start = batches->nextStart;
batches->nextStart += batches->batchSize;
const qsizetype contextStart = std::max<qsizetype>(0, start - ContextLines);
const LlmClient::BatchResult result =
co_await translateLines(job, batches->lines.mid(start, batches->batchSize), batches->lines.mid(contextStart, start - contextStart));
if (job->cancelled) {
co_return;
}
if (!result.error.isEmpty()) {
if (batches->error.isEmpty()) {
batches->error = result.error;
}
co_return;
}
for (qsizetype i = 0; i < result.translations.size(); ++i) {
batches->translations[start + i] = result.translations.at(i);
}
batches->doneLines += result.translations.size();
setProgress(double(batches->doneLines) / batches->lines.size());
setStatusText(i18n("Translating… %1 of %2 lines", batches->doneLines, batches->lines.size()));
}
}
QCoro::Task<> TranslationService::runFetchModels()
{
QPointer<TranslationService> self(this);
setFetchingModels(true);
setErrorText({});
const auto [config, configError] = co_await loadConfig();
if (!self) {
co_return;
}
if (!configError.isEmpty()) {
setErrorText(configError);
setFetchingModels(false);
co_return;
}
LlmClient client(m_network, config);
const LlmClient::ModelsResult result = co_await client.listModels();
if (!self) {
co_return;
}
if (!result.error.isEmpty()) {
setErrorText(i18n("Could not list models: %1", result.error));
} else if (result.models != m_availableModels) {
m_availableModels = result.models;
Q_EMIT availableModelsChanged();
}
setFetchingModels(false);
}
void TranslationService::fail(const std::shared_ptr<Job> &job, const QString &error)
{
if (m_job != job) {
return;
}
m_job.reset();
setErrorText(error);
setBusy(false);
setStatusText({});
}
void TranslationService::setBusy(bool busy)
{
if (busy != m_busy) {
m_busy = busy;
Q_EMIT busyChanged();
}
}
void TranslationService::setProgress(double progress)
{
if (progress != m_progress) {
m_progress = progress;
Q_EMIT progressChanged();
}
}
void TranslationService::setStatusText(const QString &text)
{
if (text != m_statusText) {
m_statusText = text;
Q_EMIT statusTextChanged();
}
}
void TranslationService::setErrorText(const QString &text)
{
if (text != m_errorText) {
m_errorText = text;
Q_EMIT errorTextChanged();
}
}
void TranslationService::setFetchingModels(bool fetching)
{
if (fetching != m_fetchingModels) {
m_fetchingModels = fetching;
Q_EMIT fetchingModelsChanged();
}
}