Skip to content

Show translation model download progress and a cancel button #1617

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
cdrini wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Show translation model download progress and a cancel button #1617

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 35 additions & 7 deletions src/plugins/translate/TranslationManager.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"};
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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<boolean>}
* @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);
}
}
114 changes: 103 additions & 11 deletions src/plugins/translate/plugin.translate.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -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")) : [];
Expand Down Expand Up @@ -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();

Expand Down Expand Up @@ -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();
}
Expand All @@ -302,6 +339,7 @@ export class TranslatePlugin extends BookReaderPlugin {
this.clearAllTranslations();
this.langToCode = e.detail.value;
this._render();
this.loadTranslationModels();
this.translateActivePageContainerElements();
}

Expand All @@ -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();
}
Expand All @@ -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();
}
}"
Expand All @@ -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"
/>`,
Expand All @@ -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;
Expand All @@ -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() {
Expand Down Expand Up @@ -435,7 +516,7 @@ export class BrTranslatePanel extends LitElement {

<div class="panel panel--start" style="text-align: right;padding: 0 10px;/*! font-size: 18px; */margin-top: 10px;">
<button class="start-translation-brn" @click="${this._toggleTranslation}" ?disabled=${!this.sourceLanguageSupported}>
${this.userTranslationActive ? "Stop Translating" : "Translate"}
${this._toggleButtonLabel()}
</button>
</div>

Expand Down Expand Up @@ -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;
}

Expand All @@ -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;
}

Expand Down Expand Up @@ -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) {
Expand All @@ -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`
<ia-activity-indicator mode="processing" style="display:block; width: 40px; height: 40px; margin: 0 auto;"></ia-activity-indicator>
<p>Downloading language model</p>
<div class="download-card">
<div class="download-card__title" title="Downloading ${pairName}">Downloading ${pairName}</div>
<div class="download-card__status">${total ? `${(total / 1e6).toFixed(1)} MB` : "\u00a0"}</div>
<div class="download-card__bar" role="progressbar" aria-label="Language model download"
aria-valuemin="0" aria-valuemax=${total} aria-valuenow=${loaded}>
<div class="download-card__bar-fill" style="width: ${total ? (loaded / total) * 100 : 0}%"></div>
</div>
</div>
`;
}
return html`<p>Language model loaded</p>`;
Expand Down
62 changes: 62 additions & 0 deletions tests/jest/plugins/translate/plugin.translate.test.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand All @@ -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);
});
});
Loading