From 449a76bd789554ec2b1bc563fc747ec81b930b8f Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 24 Oct 2026 00:58:07 +0000 Subject: [PATCH] Show translation model download progress and size Co-Authored-By: Drini Cami --- src/plugins/translate/TranslationManager.js | 42 +++++-- src/plugins/translate/plugin.translate.js | 114 ++++++++++++++++-- .../translate/plugin.translate.test.js | 62 ++++++++++ 3 files changed, 200 insertions(+), 18 deletions(-) diff --git a/src/plugins/translate/TranslationManager.js b/src/plugins/translate/TranslationManager.js index 8279a6f08..5d015a66a 100644 --- a/src/plugins/translate/TranslationManager.js +++ b/src/plugins/translate/TranslationManager.js @@ -34,6 +34,9 @@ export class TranslationManager { publicPath = ''; + /** @type {((progress: {from: string, to: string, loaded: number}) => void) | null} */ + _onModelProgress = null; + constructor() { //TODO Should default to the book language as the first element const enModel = {code: "en", name: "English (en)", type: "prod"}; @@ -64,6 +67,7 @@ export class TranslationManager { workers: 2, batchSize: 4, workerUrl: this.publicPath + '/translator-worker.js', + onmodelprogress: (progress) => this._onModelProgress?.(progress), }); const modelType = await this.translator.backing.registry; @@ -144,21 +148,45 @@ export class TranslationManager { this.currentlyTranslating[key].resolve(response.target.text); this.alreadyTranslated.add({index: key, response: response.target.text}); delete this.currentlyTranslating[key]; + }).catch((err) => { + this.currentlyTranslating[key].reject(err); + delete this.currentlyTranslating[key]; }); return promise; } /** - * Checks and updates the status of a model that is currently being retreived from IA + * Downloads the model(s) needed to translate between two languages. Pairs without + * English need two models, since they're translated through English. * @param {string} fromLanguage * @param {string} toLanguage - * @returns {Promise} + * @param {object} [options] + * @param {AbortSignal} [options.signal] + * @param {(loaded: number, total: number) => void} [options.onProgress] bytes downloaded / total */ - getTranslationModel = async(fromLanguage, toLanguage) => { - return this.translator.backing.getTranslationModel({from: fromLanguage, to: toLanguage}) - .catch((err) => { - console.error(`An error occurred while trying to retreive the ${fromLanguage}-${toLanguage} model`); - }); + loadTranslationModels = async (fromLanguage, toLanguage, {signal, onProgress} = {}) => { + const models = await this.translator.backing.getModels({from: fromLanguage, to: toLanguage}); + const total = models + .flatMap(model => Object.values(model.files)) + .reduce((sum, file) => sum + (file?.size || 0), 0); + const loadedByModel = {}; + onProgress?.(0, total); + + const handler = ({from, to, loaded}) => { + loadedByModel[`${from}${to}`] = loaded; + onProgress?.(Object.values(loadedByModel).reduce((sum, n) => sum + n, 0), total); + }; + this._onModelProgress = handler; + try { + await Promise.all(models.map(({from, to}) => this.translator.backing.getTranslationModel({from, to}, {signal}))); + } finally { + if (this._onModelProgress === handler) this._onModelProgress = null; + } + } + + /** Drops all queued translation requests */ + cancelTranslations() { + this.translator.remove(() => true); } } diff --git a/src/plugins/translate/plugin.translate.js b/src/plugins/translate/plugin.translate.js index 5037c0596..70aa8a37e 100644 --- a/src/plugins/translate/plugin.translate.js +++ b/src/plugins/translate/plugin.translate.js @@ -6,7 +6,6 @@ import { TranslationManager } from "./TranslationManager.js"; import { toISO6391, toNativeName } from '../tts/utils.js'; import { sortBy } from '../../../src/BookReader/utils.js'; import { TextSelectionManager } from '../../../src/util/TextSelectionManager.js'; -import '@internetarchive/ia-activity-indicator'; // @ts-ignore const BookReader = /** @type {typeof import('@/src/BookReader.js').default} */(window.BookReader); @@ -60,6 +59,12 @@ export class TranslatePlugin extends BookReaderPlugin { */ loadingModel = true; + /** @type {{loaded: number, total: number}} bytes of the language model(s) downloaded so far */ + downloadProgress = {loaded: 0, total: 0}; + + /** @type {AbortController | null} */ + _downloadController = null; + textSelectionManager = new TextSelectionManager('.BRtranslateLayer', this.br, {selectionElement: [".BRlineElement"]}, 1); async init() { @@ -109,6 +114,35 @@ export class TranslatePlugin extends BookReaderPlugin { this._render(); } + /** Downloads the model(s) for the current language pair, showing progress in the panel */ + async loadTranslationModels() { + this._downloadController?.abort(); + if (!this.userToggleTranslate) return; + if (this.langFromCode == this.langToCode) { + this.loadingModel = false; + this._render(); + return; + } + + const controller = this._downloadController = new AbortController(); + this.loadingModel = true; + this.downloadProgress = {loaded: 0, total: 0}; + this._render(); + try { + await this.translationManager.loadTranslationModels(this.langFromCode, this.langToCode, { + signal: controller.signal, + onProgress: (loaded, total) => { + this.downloadProgress = {loaded, total}; + if (this._panel) this._panel.downloadProgress = this.downloadProgress; + }, + }); + this.loadingModel = false; + this._render(); + } catch (err) { + if (!controller.signal.aborted) console.error(`Could not load the ${this.langFromCode}-${this.langToCode} translation model`, err); + } + } + /** @param {HTMLElement} page*/ getParagraphsOnPage = (page) => { return page ? Array.from(page.querySelectorAll(".BRtextLayer > .BRparagraphElement")) : []; @@ -189,11 +223,13 @@ export class TranslatePlugin extends BookReaderPlugin { if (paragraph.textContent.length !== 0) { const pagePriority = parseFloat(pageIndex) + priority + pidx; - this.translationManager.getTranslationModel(this.langFromCode, this.langToCode).then(() => { - this._panel.loadingModel = false; - this.loadingModel = false; - }); - const translatedText = await this.translationManager.getTranslation(this.langFromCode, this.langToCode, pageIndex, pidx, paragraph.textContent, pagePriority); + let translatedText; + try { + translatedText = await this.translationManager.getTranslation(this.langFromCode, this.langToCode, pageIndex, pidx, paragraph.textContent, pagePriority); + } catch { + // Cancelled, e.g. translation was turned off while waiting + return; + } // prevent duplicate spans from appearing if exists translatedParagraph.firstElementChild?.remove(); @@ -293,6 +329,7 @@ export class TranslatePlugin extends BookReaderPlugin { console.log(this.langFromCode, this.langToCode); this._render(); + this.loadTranslationModels(); if (this.langFromCode !== this.langToCode) { this.translateActivePageContainerElements(); } @@ -302,6 +339,7 @@ export class TranslatePlugin extends BookReaderPlugin { this.clearAllTranslations(); this.langToCode = e.detail.value; this._render(); + this.loadTranslationModels(); this.translateActivePageContainerElements(); } @@ -319,11 +357,14 @@ export class TranslatePlugin extends BookReaderPlugin { this._render(); if (!this.userToggleTranslate) { + this._downloadController?.abort(); + this.translationManager.cancelTranslations(); this.clearAllTranslations(); this.br.trigger('translationDisabled', { }); this.textSelectionManager.detach(); } else { this.br.trigger('translationEnabled', { }); + this.loadTranslationModels(); this.translateActivePageContainerElements(); this.textSelectionManager.attach(); } @@ -350,6 +391,7 @@ export class TranslatePlugin extends BookReaderPlugin { this._panel.detectedToLang = this.langToCode; this._panel.detectedFromLang = this.langFromCode; this._panel.loadingModel = this.loadingModel; + this._panel.downloadProgress = this.downloadProgress; this._panel.sourceLanguageSupported = this.isSourceLanguageSupported(); } }" @@ -363,6 +405,7 @@ export class TranslatePlugin extends BookReaderPlugin { .detectedFromLang=${this.langFromCode} .detectedToLang=${this.langToCode} .loadingModel=${this.loadingModel} + .downloadProgress=${this.downloadProgress} .sourceLanguageSupported=${this.isSourceLanguageSupported()} class="translate-panel" />`, @@ -382,6 +425,7 @@ export class BrTranslatePanel extends LitElement { @property({ type: String }) detectedFromLang = ''; @property({ type: String }) detectedToLang = ''; @property({ type: Boolean }) loadingModel; + @property({ type: Object }) downloadProgress = {loaded: 0, total: 0}; // bytes @property({ type: Boolean }) sourceLanguageSupported = true; @query('#lang-from') langFromDropdown; @@ -402,6 +446,43 @@ export class BrTranslatePanel extends LitElement { background-color: rgba(255, 186, 8, 0.18); color: #ffe7a3; } + + .lang-models-loading { + padding: 0 10px; + text-align: center; + } + + .download-card { + padding: 12px; + border-radius: 8px; + background-color: rgba(255, 255, 255, 0.08); + text-align: left; + } + + .download-card__title { + font-size: 13px; + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; + } + + .download-card__status { + font-size: 12px; + color: rgba(255, 255, 255, 0.6); + } + + .download-card__bar { + height: 4px; + margin-top: 12px; + border-radius: 2px; + overflow: hidden; + background-color: rgba(255, 255, 255, 0.15); + } + + .download-card__bar-fill { + height: 100%; + background-color: #6fb3ff; + } `; connectedCallback() { @@ -435,7 +516,7 @@ export class BrTranslatePanel extends LitElement {
@@ -482,7 +563,6 @@ export class BrTranslatePanel extends LitElement { if (this._getSelectedLang('to') !== this._getSelectedLang('from')) { this.prevSelectedLang = this._getSelectedLang('from'); } - this.loadingModel = true; this.detectedFromLang = event.target.value; } @@ -498,7 +578,6 @@ export class BrTranslatePanel extends LitElement { if (this._getSelectedLang('from') !== event.target.value) { this.prevSelectedLang = this._getSelectedLang('from'); } - this.loadingModel = true; this.detectedToLang = event.target.value; } @@ -533,6 +612,11 @@ export class BrTranslatePanel extends LitElement { this.dispatchEvent(toggleTranslateEvent); } + _toggleButtonLabel() { + if (!this.userTranslationActive) return "Translate"; + return this.loadingModel ? "Cancel" : "Stop Translating"; + } + // TODO: Hardcoded warning message for now but should add more statuses _statusWarning() { if (!this.sourceLanguageSupported) { @@ -547,9 +631,17 @@ export class BrTranslatePanel extends LitElement { _languageModelStatus() { if (this.userTranslationActive) { if (this.loadingModel) { + const {loaded, total} = this.downloadProgress; + const pairName = `${toNativeName(this.detectedFromLang) || this.detectedFromLang} → ${toNativeName(this.detectedToLang) || this.detectedToLang}`; return html` - -

Downloading language model

+
+
Downloading ${pairName}
+
${total ? `${(total / 1e6).toFixed(1)} MB` : "\u00a0"}
+
+
+
+
`; } return html`

Language model loaded

`; diff --git a/tests/jest/plugins/translate/plugin.translate.test.js b/tests/jest/plugins/translate/plugin.translate.test.js index 81dbd5159..be7fa3d60 100644 --- a/tests/jest/plugins/translate/plugin.translate.test.js +++ b/tests/jest/plugins/translate/plugin.translate.test.js @@ -4,7 +4,9 @@ jest.mock('@internetarchive/bergamot-translator/translator.js', () => ({ BatchTranslator: class {}, })); +import { render } from 'lit'; import { BrTranslatePanel } from '@/src/plugins/translate/plugin.translate.js'; +import { TranslationManager } from '@/src/plugins/translate/TranslationManager.js'; describe('BrTranslatePanel unsupported language warning', () => { test('shows unsupported language warning message', () => { @@ -15,3 +17,63 @@ describe('BrTranslatePanel unsupported language warning', () => { expect(panel._statusWarning()).toBe('Translating from zzzz is not supported'); }); }); + +describe('BrTranslatePanel download progress', () => { + test('shows the language pair, total size and progress while loading a model', () => { + const panel = new BrTranslatePanel(); + panel.userTranslationActive = true; + panel.loadingModel = true; + panel.detectedFromLang = 'en'; + panel.detectedToLang = 'fr'; + panel.downloadProgress = {loaded: 12_950_000, total: 25_900_000}; + + const container = document.createElement('div'); + render(panel._languageModelStatus(), container); + expect(container.querySelector('.download-card__title').textContent).toBe('Downloading English → français'); + expect(container.querySelector('.download-card__status').textContent).toBe('25.9 MB'); + expect(container.querySelector('[role=progressbar]').getAttribute('aria-valuenow')).toBe('12950000'); + expect(container.querySelector('.download-card__bar-fill').style.width).toBe('50%'); + }); + + test('leaves the size blank until the total is known', () => { + const panel = new BrTranslatePanel(); + panel.userTranslationActive = true; + panel.loadingModel = true; + panel.downloadProgress = {loaded: 0, total: 0}; + + const container = document.createElement('div'); + render(panel._languageModelStatus(), container); + expect(container.querySelector('.download-card__status').textContent).toBe('\u00a0'); + }); + + test('toggle button reads Cancel while a model is downloading', () => { + const panel = new BrTranslatePanel(); + expect(panel._toggleButtonLabel()).toBe('Translate'); + panel.userTranslationActive = true; + panel.loadingModel = true; + expect(panel._toggleButtonLabel()).toBe('Cancel'); + panel.loadingModel = false; + expect(panel._toggleButtonLabel()).toBe('Stop Translating'); + }); +}); + +describe('TranslationManager.loadTranslationModels', () => { + test('reports combined progress across pivot models', async () => { + const manager = new TranslationManager(); + const models = [ + {from: 'fr', to: 'en', files: {model: {size: 30}, vocab: {size: 10}}}, + {from: 'en', to: 'de', files: {model: {size: 50}, lex: {size: 10}}}, + ]; + manager.translator = /** @type {any} */ ({ + backing: { + getModels: async () => models, + getTranslationModel: async ({from, to}) => manager._onModelProgress({from, to, loaded: 20}), + }, + }); + + const onProgress = jest.fn(); + await manager.loadTranslationModels('fr', 'de', {onProgress}); + expect(onProgress).toHaveBeenNthCalledWith(1, 0, 100); + expect(onProgress).toHaveBeenLastCalledWith(40, 100); + }); +});