nllw 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- nllw-0.1.0/.gitignore +72 -0
- nllw-0.1.0/PKG-INFO +94 -0
- nllw-0.1.0/README.md +65 -0
- nllw-0.1.0/architecture_NLLW.png +0 -0
- nllw-0.1.0/compare_translation_approaches.py +129 -0
- nllw-0.1.0/demo.gif +0 -0
- nllw-0.1.0/french_to_english.png +0 -0
- nllw-0.1.0/nllw/__init__.py +42 -0
- nllw-0.1.0/nllw/core.py +252 -0
- nllw-0.1.0/nllw/languages.py +275 -0
- nllw-0.1.0/nllw/translation.py +199 -0
- nllw-0.1.0/nllw.egg-info/PKG-INFO +94 -0
- nllw-0.1.0/nllw.egg-info/SOURCES.txt +18 -0
- nllw-0.1.0/nllw.egg-info/dependency_links.txt +1 -0
- nllw-0.1.0/nllw.egg-info/requires.txt +8 -0
- nllw-0.1.0/nllw.egg-info/top_level.txt +1 -0
- nllw-0.1.0/pyproject.toml +59 -0
- nllw-0.1.0/setup.cfg +4 -0
- nllw-0.1.0/supported_languages.md +210 -0
- nllw-0.1.0/textual_interface.py +250 -0
nllw-0.1.0/.gitignore
ADDED
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
# Python
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[cod]
|
|
4
|
+
*$py.class
|
|
5
|
+
*.so
|
|
6
|
+
.Python
|
|
7
|
+
|
|
8
|
+
# Distribution / packaging
|
|
9
|
+
build/
|
|
10
|
+
dist/
|
|
11
|
+
*.egg-info/
|
|
12
|
+
*.egg
|
|
13
|
+
MANIFEST
|
|
14
|
+
|
|
15
|
+
# PyInstaller
|
|
16
|
+
*.manifest
|
|
17
|
+
*.spec
|
|
18
|
+
|
|
19
|
+
# Installer logs
|
|
20
|
+
pip-log.txt
|
|
21
|
+
pip-delete-this-directory.txt
|
|
22
|
+
|
|
23
|
+
# Unit test / coverage
|
|
24
|
+
htmlcov/
|
|
25
|
+
.tox/
|
|
26
|
+
.nox/
|
|
27
|
+
.coverage
|
|
28
|
+
.coverage.*
|
|
29
|
+
.cache
|
|
30
|
+
nosetests.xml
|
|
31
|
+
coverage.xml
|
|
32
|
+
*.cover
|
|
33
|
+
.hypothesis/
|
|
34
|
+
.pytest_cache/
|
|
35
|
+
other/*
|
|
36
|
+
|
|
37
|
+
# Jupyter Notebook
|
|
38
|
+
.ipynb_checkpoints
|
|
39
|
+
*.ipynb
|
|
40
|
+
|
|
41
|
+
# pyenv
|
|
42
|
+
.python-version
|
|
43
|
+
|
|
44
|
+
# Environments
|
|
45
|
+
.env
|
|
46
|
+
.venv
|
|
47
|
+
env/
|
|
48
|
+
venv/
|
|
49
|
+
ENV/
|
|
50
|
+
env.bak/
|
|
51
|
+
venv.bak/
|
|
52
|
+
|
|
53
|
+
# IDEs
|
|
54
|
+
.vscode/
|
|
55
|
+
.idea/
|
|
56
|
+
*.swp
|
|
57
|
+
*.swo
|
|
58
|
+
*~
|
|
59
|
+
|
|
60
|
+
# OS
|
|
61
|
+
.DS_Store
|
|
62
|
+
Thumbs.db
|
|
63
|
+
|
|
64
|
+
# Project specific
|
|
65
|
+
transformers-main/*
|
|
66
|
+
docs_hf/*
|
|
67
|
+
nllb-200-distilled-*-ctranslate2/
|
|
68
|
+
*.pkl
|
|
69
|
+
*.csv
|
|
70
|
+
*.png
|
|
71
|
+
*.gif
|
|
72
|
+
|
nllw-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,94 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: nllw
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Simultaneous Machine Translation (SimulMT) with NLLB model optimization
|
|
5
|
+
Author: Quentin Fuxa
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/QuentinFuxa/NoLanguageLeftWaiting
|
|
8
|
+
Keywords: machine-translation,simultaneous-translation,nllb,streaming,nlp,transformers,ai
|
|
9
|
+
Classifier: Development Status :: 4 - Beta
|
|
10
|
+
Classifier: Intended Audience :: Developers
|
|
11
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.15
|
|
19
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
20
|
+
Classifier: Topic :: Multimedia :: Sound/Audio :: Speech
|
|
21
|
+
Requires-Python: >=3.8
|
|
22
|
+
Description-Content-Type: text/markdown
|
|
23
|
+
Requires-Dist: torch>=1.10.0
|
|
24
|
+
Requires-Dist: transformers>=4.30.0
|
|
25
|
+
Provides-Extra: ctranslate2
|
|
26
|
+
Requires-Dist: ctranslate2>=3.16.0; extra == "ctranslate2"
|
|
27
|
+
Provides-Extra: demo
|
|
28
|
+
Requires-Dist: textual>=0.40.0; extra == "demo"
|
|
29
|
+
|
|
30
|
+
<h1 align="center">NoLanguageLeftWaiting</h1>
|
|
31
|
+
|
|
32
|
+
<p align="center">
|
|
33
|
+
<img src="demo.gif"width="730">
|
|
34
|
+
</p>
|
|
35
|
+
|
|
36
|
+
<p align="center">
|
|
37
|
+
<img src="architecture_NLLW.png"width="730">
|
|
38
|
+
</p>
|
|
39
|
+
|
|
40
|
+
Converts [NoLanguageLeftBehind](https://arxiv.org/abs/2207.04672) translation model to a SimulMT (Simultaneous Machine Translation) model, optimized for live/streaming use cases.
|
|
41
|
+
|
|
42
|
+
> Based offline models such as NLLB suffer from eos token and punctuation insertion, inconsistent prefix handling and exponentially growing computational overhead as input length increases. This implementation aims at resolving that.
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
- [LocalAgreement policy](https://www.isca-archive.org/interspeech_2020/liu20s_interspeech.pdf)
|
|
46
|
+
- [HuggingFace transformers](https://huggingface.co/docs/transformers/model_doc/auto#transformers.AutoModelForSeq2SeqLM) implementation only.
|
|
47
|
+
- Built for [WhisperLiveKit](https://github.com/QuentinFuxa/WhisperLiveKit)
|
|
48
|
+
- 200 languages. See [supported_languages.md](supported_languages.md) for the full list.
|
|
49
|
+
|
|
50
|
+
## Installation
|
|
51
|
+
|
|
52
|
+
```bash
|
|
53
|
+
pip install nllw
|
|
54
|
+
```
|
|
55
|
+
> The textual frontend is not installed by default.
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
## Quick Start
|
|
59
|
+
|
|
60
|
+
1. Demo interface :
|
|
61
|
+
```bash
|
|
62
|
+
python textual_interface.py
|
|
63
|
+
```
|
|
64
|
+
|
|
65
|
+
2. Use it as a package
|
|
66
|
+
```python
|
|
67
|
+
import nllw
|
|
68
|
+
|
|
69
|
+
model = nllw.load_model(
|
|
70
|
+
src_langs=["fra_Latn"],
|
|
71
|
+
nllb_backend="transformers",
|
|
72
|
+
nllb_size="600M"
|
|
73
|
+
)
|
|
74
|
+
translator = nllw.OnlineTranslation(
|
|
75
|
+
model,
|
|
76
|
+
input_languages=["fra_Latn"],
|
|
77
|
+
output_languages=["eng_Latn"]
|
|
78
|
+
)
|
|
79
|
+
|
|
80
|
+
translator.insert_tokens(tokens)
|
|
81
|
+
validated, buffer = translator.process()
|
|
82
|
+
|
|
83
|
+
print(f"Stable: {validated[0].text}")
|
|
84
|
+
print(f"Buffer: {buffer.text}")
|
|
85
|
+
```
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
## Input vs Output length:
|
|
89
|
+
|
|
90
|
+
Succesfully maintain output length, even if stable prefix tends to take time to grow.
|
|
91
|
+
|
|
92
|
+
<p align="center">
|
|
93
|
+
<img src="french_to_english.png"width="730">
|
|
94
|
+
</p>
|
nllw-0.1.0/README.md
ADDED
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
<h1 align="center">NoLanguageLeftWaiting</h1>
|
|
2
|
+
|
|
3
|
+
<p align="center">
|
|
4
|
+
<img src="demo.gif"width="730">
|
|
5
|
+
</p>
|
|
6
|
+
|
|
7
|
+
<p align="center">
|
|
8
|
+
<img src="architecture_NLLW.png"width="730">
|
|
9
|
+
</p>
|
|
10
|
+
|
|
11
|
+
Converts [NoLanguageLeftBehind](https://arxiv.org/abs/2207.04672) translation model to a SimulMT (Simultaneous Machine Translation) model, optimized for live/streaming use cases.
|
|
12
|
+
|
|
13
|
+
> Based offline models such as NLLB suffer from eos token and punctuation insertion, inconsistent prefix handling and exponentially growing computational overhead as input length increases. This implementation aims at resolving that.
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
- [LocalAgreement policy](https://www.isca-archive.org/interspeech_2020/liu20s_interspeech.pdf)
|
|
17
|
+
- [HuggingFace transformers](https://huggingface.co/docs/transformers/model_doc/auto#transformers.AutoModelForSeq2SeqLM) implementation only.
|
|
18
|
+
- Built for [WhisperLiveKit](https://github.com/QuentinFuxa/WhisperLiveKit)
|
|
19
|
+
- 200 languages. See [supported_languages.md](supported_languages.md) for the full list.
|
|
20
|
+
|
|
21
|
+
## Installation
|
|
22
|
+
|
|
23
|
+
```bash
|
|
24
|
+
pip install nllw
|
|
25
|
+
```
|
|
26
|
+
> The textual frontend is not installed by default.
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
## Quick Start
|
|
30
|
+
|
|
31
|
+
1. Demo interface :
|
|
32
|
+
```bash
|
|
33
|
+
python textual_interface.py
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
2. Use it as a package
|
|
37
|
+
```python
|
|
38
|
+
import nllw
|
|
39
|
+
|
|
40
|
+
model = nllw.load_model(
|
|
41
|
+
src_langs=["fra_Latn"],
|
|
42
|
+
nllb_backend="transformers",
|
|
43
|
+
nllb_size="600M"
|
|
44
|
+
)
|
|
45
|
+
translator = nllw.OnlineTranslation(
|
|
46
|
+
model,
|
|
47
|
+
input_languages=["fra_Latn"],
|
|
48
|
+
output_languages=["eng_Latn"]
|
|
49
|
+
)
|
|
50
|
+
|
|
51
|
+
translator.insert_tokens(tokens)
|
|
52
|
+
validated, buffer = translator.process()
|
|
53
|
+
|
|
54
|
+
print(f"Stable: {validated[0].text}")
|
|
55
|
+
print(f"Buffer: {buffer.text}")
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
## Input vs Output length:
|
|
60
|
+
|
|
61
|
+
Succesfully maintain output length, even if stable prefix tends to take time to grow.
|
|
62
|
+
|
|
63
|
+
<p align="center">
|
|
64
|
+
<img src="french_to_english.png"width="730">
|
|
65
|
+
</p>
|
|
Binary file
|
|
@@ -0,0 +1,129 @@
|
|
|
1
|
+
from translation_backend import TranslationBackend
|
|
2
|
+
import matplotlib.pyplot as plt
|
|
3
|
+
import pandas as pd
|
|
4
|
+
|
|
5
|
+
src_texts = [
|
|
6
|
+
"Il s'arrêtait par moments devant une villa",
|
|
7
|
+
"coquettement nichée dans la verdure, il regardait",
|
|
8
|
+
"par la grille et voyait au loin des femmes",
|
|
9
|
+
"élégantes sur les balcons et les terrasses, des",
|
|
10
|
+
"enfants couraient dans les jardins. Il s'intéressait",
|
|
11
|
+
"surtout aux fleurs ; c'étaient elles qui attiraient",
|
|
12
|
+
"particulièrement ses regards. De temps en temps,",
|
|
13
|
+
"il voyait passer des cavaliers, des amazones et de",
|
|
14
|
+
"belles voitures ; il les suivait d'un œil curieux et",
|
|
15
|
+
"les oubliait avant qu'ils eussent disparu. ",
|
|
16
|
+
"Une fois, il s'arrêta et compta son argent ; il",
|
|
17
|
+
"lui restait trente kopecks : « vingt au sergent de",
|
|
18
|
+
"ville, trois à Nastassia pour la lettre, j'en ai donc",
|
|
19
|
+
"donné hier à Marmeladov quarante-sept ou",
|
|
20
|
+
"cinquante », se dit-il. Il devait avoir une raison de",
|
|
21
|
+
"calculer ainsi, mais il l'oublia en tirant l'argent de",
|
|
22
|
+
"sa poche et ne s'en souvint qu'un peu plus tard en",
|
|
23
|
+
"passant devant un marchand de comestibles, une",
|
|
24
|
+
"sorte de gargote plutôt ; il sentit alors qu'il avait",
|
|
25
|
+
"faim. "
|
|
26
|
+
]
|
|
27
|
+
|
|
28
|
+
# print("1: direct generation)")
|
|
29
|
+
# translation_backend_1 = TranslationBackend(source_lang='fra_Latn', target_lang="eng_Latn")
|
|
30
|
+
# l_vals_no_cache = []
|
|
31
|
+
|
|
32
|
+
# for i in range(1, len(src_texts) + 1):
|
|
33
|
+
# print(f'{i}/{len(src_texts) + 1}')
|
|
34
|
+
# truncated_text = " ".join(src_texts[:i])
|
|
35
|
+
# input_tokens = translation_backend_1.tokenizer(truncated_text, return_tensors="pt").to(translation_backend_1.device)
|
|
36
|
+
# encoder_outputs = translation_backend_1.model.get_encoder()(**input_tokens)
|
|
37
|
+
# output_tokens = translation_backend_1.model.generate(
|
|
38
|
+
# encoder_outputs=encoder_outputs,
|
|
39
|
+
# forced_bos_token_id=translation_backend_1.bos_token_id
|
|
40
|
+
# )
|
|
41
|
+
# output_text = translation_backend_1.tokenizer.decode(output_tokens[0], skip_special_tokens=True)
|
|
42
|
+
# l_vals_no_cache.append({
|
|
43
|
+
# "input": truncated_text,
|
|
44
|
+
# "output_text": output_text,
|
|
45
|
+
# "input_word_count": len(truncated_text.split()),
|
|
46
|
+
# "output_word_count": len(output_text.split())
|
|
47
|
+
# })
|
|
48
|
+
|
|
49
|
+
# print("2: prefix reuse")
|
|
50
|
+
# translation_backend_2 = TranslationBackend(source_lang='fra_Latn', target_lang="eng_Latn")
|
|
51
|
+
# l_vals_with_cache = []
|
|
52
|
+
|
|
53
|
+
# for i in range(1, len(src_texts) + 1):
|
|
54
|
+
# print(f'{i}/{len(src_texts) + 1}')
|
|
55
|
+
# truncated_text = " ".join(src_texts[:i])
|
|
56
|
+
# stable_translation, buffer = translation_backend_2.translate(truncated_text)
|
|
57
|
+
|
|
58
|
+
# full_output = stable_translation + buffer
|
|
59
|
+
# l_vals_with_cache.append({
|
|
60
|
+
# "input": truncated_text,
|
|
61
|
+
# "stable_translation": stable_translation,
|
|
62
|
+
# "buffer": buffer,
|
|
63
|
+
# "full_output": full_output,
|
|
64
|
+
# "input_word_count": len(truncated_text.split()),
|
|
65
|
+
# "stable_word_count": len(stable_translation.split()) if stable_translation else 0,
|
|
66
|
+
# "buffer_word_count": len(buffer.split()) if buffer else 0,
|
|
67
|
+
# "total_output_word_count": len(full_output.split()) if full_output else 0
|
|
68
|
+
# })
|
|
69
|
+
|
|
70
|
+
# df_base = pd.DataFrame(l_vals_no_cache)
|
|
71
|
+
# df_prefix = pd.DataFrame(l_vals_with_cache)
|
|
72
|
+
# df_base.to_csv('output_base_analysis.csv')
|
|
73
|
+
# df_prefix.to_csv('output_prefix_analysis.csv')
|
|
74
|
+
|
|
75
|
+
df_base = pd.read_csv('output_base_analysis.csv')
|
|
76
|
+
df_prefix = pd.read_csv('output_prefix_analysis.csv')
|
|
77
|
+
|
|
78
|
+
# df_merged = pd.concat([df_base, df_prefix], axis=1)
|
|
79
|
+
|
|
80
|
+
# Extract data from dataframe
|
|
81
|
+
iterations = list(range(1, len(df_prefix) + 1))
|
|
82
|
+
input_counts = df_prefix['input_word_count'].tolist()
|
|
83
|
+
output_with_prefix_and_buffer = df_prefix['total_output_word_count'].tolist()
|
|
84
|
+
stable_counts = df_prefix['stable_word_count'].tolist()
|
|
85
|
+
buffer_counts = df_prefix['buffer_word_count'].tolist()
|
|
86
|
+
base_output_wc = df_base['output_word_count'].tolist()
|
|
87
|
+
|
|
88
|
+
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 10))
|
|
89
|
+
fig, (ax1) = plt.subplots(1, 1, figsize=(5, 5))
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
colors = {
|
|
93
|
+
"dark_gray_1": "#2E2E2E", # deep graphite gray
|
|
94
|
+
"dark_gray_2": "#4B4B4B", # medium-dark neutral gray
|
|
95
|
+
"orange": "#FF7F0E", # modern vivid orange (Matplotlib's default orange)
|
|
96
|
+
"red": "#D62728", # elegant deep red
|
|
97
|
+
}
|
|
98
|
+
from matplotlib.ticker import MaxNLocator
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
ax1.plot(iterations, input_counts, label='Input', marker='o', color=colors['dark_gray_1'])
|
|
102
|
+
ax1.plot(iterations, base_output_wc, label='Base output (no prefix, no stability)', marker='^', color=colors['dark_gray_2'])
|
|
103
|
+
ax1.plot(iterations, stable_counts, label='With prefix (Stable part)', marker='^', linewidth=2, color=colors['red'])
|
|
104
|
+
ax1.plot(iterations, output_with_prefix_and_buffer, label='With prefix use output (Stable part + Buffer)', marker='^', color=colors['orange'])
|
|
105
|
+
ax1.set_xlabel('Iteration')
|
|
106
|
+
ax1.set_ylabel('Word count')
|
|
107
|
+
ax1.set_title('SimulMT french to english Word Count')
|
|
108
|
+
ax1.legend(fontsize=10)
|
|
109
|
+
ax1.grid(True, alpha=0.3)
|
|
110
|
+
ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
|
|
111
|
+
|
|
112
|
+
# # Second plot: Detailed view of cached approach
|
|
113
|
+
# ax2.plot(iterations, input_counts, label='Input', marker='o', linewidth=2)
|
|
114
|
+
# ax2.plot(iterations, stable_counts, label='Stable prefix (cached)', marker='D', linewidth=2)
|
|
115
|
+
# ax2.plot(iterations, buffer_counts, label='Buffer (new)', marker='x', linewidth=2)
|
|
116
|
+
# ax2.plot(iterations, output_with_cache, label='Total output', marker='^', linewidth=2, linestyle='--')
|
|
117
|
+
# ax2.set_xlabel('Iteration', fontsize=12)
|
|
118
|
+
# ax2.set_ylabel('Word count', fontsize=12)
|
|
119
|
+
# ax2.set_title('Detailed view: Stable prefix vs Buffer (with cache)', fontsize=14, fontweight='bold')
|
|
120
|
+
# ax2.legend(fontsize=10)
|
|
121
|
+
# ax2.grid(True, alpha=0.3)
|
|
122
|
+
|
|
123
|
+
plt.tight_layout()
|
|
124
|
+
plt.savefig('translation_comparison.png', dpi=300, bbox_inches='tight')
|
|
125
|
+
plt.show()
|
|
126
|
+
|
|
127
|
+
# df = pd.DataFrame(l_vals_with_cache)
|
|
128
|
+
# df.to_csv('output_prefix_analysis.csv')
|
|
129
|
+
# print("terminated")
|
nllw-0.1.0/demo.gif
ADDED
|
Binary file
|
|
Binary file
|
|
@@ -0,0 +1,42 @@
|
|
|
1
|
+
from .translation import (
|
|
2
|
+
load_model,
|
|
3
|
+
OnlineTranslation,
|
|
4
|
+
TranslationModel,
|
|
5
|
+
TimedText,
|
|
6
|
+
MIN_SILENCE_DURATION_DEL_BUFFER,
|
|
7
|
+
)
|
|
8
|
+
|
|
9
|
+
from .core import TranslationBackend
|
|
10
|
+
|
|
11
|
+
from .languages import (
|
|
12
|
+
get_nllb_code,
|
|
13
|
+
get_language_code_code,
|
|
14
|
+
get_language_name_by_language_code,
|
|
15
|
+
get_language_name_by_nllb,
|
|
16
|
+
get_language_info,
|
|
17
|
+
list_all_languages,
|
|
18
|
+
list_all_nllb_codes,
|
|
19
|
+
list_all_language_code_codes,
|
|
20
|
+
LANGUAGES,
|
|
21
|
+
)
|
|
22
|
+
|
|
23
|
+
__all__ = [
|
|
24
|
+
# Main API
|
|
25
|
+
"load_model",
|
|
26
|
+
"OnlineTranslation",
|
|
27
|
+
"TranslationModel",
|
|
28
|
+
"TimedText",
|
|
29
|
+
"MIN_SILENCE_DURATION_DEL_BUFFER",
|
|
30
|
+
# Backend (advanced)
|
|
31
|
+
"TranslationBackend",
|
|
32
|
+
# Language utilities
|
|
33
|
+
"get_nllb_code",
|
|
34
|
+
"get_language_code_code",
|
|
35
|
+
"get_language_name_by_language_code",
|
|
36
|
+
"get_language_name_by_nllb",
|
|
37
|
+
"get_language_info",
|
|
38
|
+
"list_all_languages",
|
|
39
|
+
"list_all_nllb_codes",
|
|
40
|
+
"list_all_language_code_codes",
|
|
41
|
+
"LANGUAGES",
|
|
42
|
+
]
|
nllw-0.1.0/nllw/core.py
ADDED
|
@@ -0,0 +1,252 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
|
|
3
|
+
from transformers.cache_utils import EncoderDecoderCache, DynamicCache
|
|
4
|
+
from typing import Tuple, Optional
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
class TranslationBackend:
|
|
8
|
+
def __init__(self, source_lang, target_lang, model_name: str = "facebook/nllb-200-distilled-600M", model=None, tokenizer=None):
|
|
9
|
+
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
10
|
+
|
|
11
|
+
if model is not None:
|
|
12
|
+
self.model = model
|
|
13
|
+
if not hasattr(model, 'device') or str(model.device) != self.device:
|
|
14
|
+
self.model = self.model.to(self.device)
|
|
15
|
+
else:
|
|
16
|
+
self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(self.device)
|
|
17
|
+
|
|
18
|
+
if tokenizer is not None:
|
|
19
|
+
self.tokenizer = tokenizer
|
|
20
|
+
else:
|
|
21
|
+
self.tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang=source_lang)
|
|
22
|
+
|
|
23
|
+
self.source_lang = source_lang
|
|
24
|
+
self.target_lang = target_lang
|
|
25
|
+
self.bos_token_id = self.tokenizer.convert_tokens_to_ids(self.target_lang)
|
|
26
|
+
|
|
27
|
+
self.sentence_end_token_ids = set()
|
|
28
|
+
|
|
29
|
+
# find all tokens that decode to sentence-ending punctuation. For ex: token 81 and 248075 both represent '.')
|
|
30
|
+
for token_id in range(min(300000, self.tokenizer.vocab_size)):
|
|
31
|
+
try:
|
|
32
|
+
decoded = self.tokenizer.decode([token_id])
|
|
33
|
+
cleaned = decoded.strip().strip("'\"").strip()
|
|
34
|
+
if cleaned in ['.', '!', '?']:
|
|
35
|
+
self.sentence_end_token_ids.add(token_id)
|
|
36
|
+
except:
|
|
37
|
+
pass
|
|
38
|
+
|
|
39
|
+
self.previous_tokens = None
|
|
40
|
+
self.stable_prefix = None
|
|
41
|
+
|
|
42
|
+
def simple_translation(self, text):
|
|
43
|
+
inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
|
|
44
|
+
encoder_outputs = self.model.get_encoder()(**inputs)
|
|
45
|
+
output = self.generate(
|
|
46
|
+
encoder_outputs=encoder_outputs,
|
|
47
|
+
)
|
|
48
|
+
result = self.tokenizer.decode(output[0], skip_special_tokens=True)
|
|
49
|
+
return output, result
|
|
50
|
+
|
|
51
|
+
def generate(
|
|
52
|
+
self,
|
|
53
|
+
encoder_outputs: torch.Tensor,
|
|
54
|
+
attention_mask: Optional[torch.Tensor] = None,
|
|
55
|
+
max_length: Optional[int] = 200
|
|
56
|
+
|
|
57
|
+
) -> Tuple[torch.Tensor, Optional[Tuple]]:
|
|
58
|
+
with torch.no_grad():
|
|
59
|
+
generated_tokens = self.model.generate(
|
|
60
|
+
encoder_outputs=encoder_outputs,
|
|
61
|
+
attention_mask=attention_mask,
|
|
62
|
+
forced_bos_token_id=self.bos_token_id,
|
|
63
|
+
max_length=max_length,
|
|
64
|
+
)
|
|
65
|
+
return generated_tokens
|
|
66
|
+
|
|
67
|
+
def compute_common_prefix_tokens(
|
|
68
|
+
self, new_tokens
|
|
69
|
+
):
|
|
70
|
+
common_length = 0
|
|
71
|
+
for i in range(min(len(self.previous_tokens[0]), len(new_tokens[0]))):
|
|
72
|
+
if self.previous_tokens[0][i] != new_tokens[0][i]:
|
|
73
|
+
common_length = i
|
|
74
|
+
break
|
|
75
|
+
else:
|
|
76
|
+
common_length = min(len(self.previous_tokens[0]), len(new_tokens[0]))
|
|
77
|
+
|
|
78
|
+
last_sentence_end = -1
|
|
79
|
+
for i in range(common_length):
|
|
80
|
+
if new_tokens[0][i].item() in self.sentence_end_token_ids:
|
|
81
|
+
last_sentence_end = i
|
|
82
|
+
|
|
83
|
+
if last_sentence_end >= 0:
|
|
84
|
+
return new_tokens[:, :last_sentence_end]
|
|
85
|
+
|
|
86
|
+
return new_tokens[:, :common_length]
|
|
87
|
+
|
|
88
|
+
def translate(self, text: str) -> str:
|
|
89
|
+
word_count = len(text.strip().split())
|
|
90
|
+
|
|
91
|
+
if word_count < 3:
|
|
92
|
+
return "", ""
|
|
93
|
+
|
|
94
|
+
inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
|
|
95
|
+
with torch.no_grad():
|
|
96
|
+
encoder_outputs = self.model.get_encoder()(**inputs)
|
|
97
|
+
|
|
98
|
+
if (self.previous_tokens is not None and self.stable_prefix is not None):
|
|
99
|
+
translation_tokens = self._continue_generation_with_cache(
|
|
100
|
+
encoder_hidden_states=encoder_outputs.last_hidden_state,
|
|
101
|
+
)
|
|
102
|
+
else:
|
|
103
|
+
with torch.no_grad():
|
|
104
|
+
translation_tokens = self.generate(
|
|
105
|
+
encoder_outputs=encoder_outputs,
|
|
106
|
+
attention_mask=inputs['attention_mask'],
|
|
107
|
+
)
|
|
108
|
+
if self.previous_tokens is not None:
|
|
109
|
+
self.stable_prefix = self.compute_common_prefix_tokens(new_tokens=translation_tokens)
|
|
110
|
+
self.previous_tokens = translation_tokens
|
|
111
|
+
|
|
112
|
+
buffer = self.tokenizer.decode(
|
|
113
|
+
translation_tokens[0],
|
|
114
|
+
skip_special_tokens=True
|
|
115
|
+
)
|
|
116
|
+
if self.stable_prefix is not None:
|
|
117
|
+
stable_translation = self.tokenizer.decode(
|
|
118
|
+
self.stable_prefix[0],
|
|
119
|
+
skip_special_tokens=True
|
|
120
|
+
)
|
|
121
|
+
return stable_translation, buffer[len(stable_translation):]
|
|
122
|
+
else:
|
|
123
|
+
return "", buffer
|
|
124
|
+
|
|
125
|
+
def _continue_generation_with_cache(
|
|
126
|
+
self,
|
|
127
|
+
encoder_hidden_states: torch.Tensor,
|
|
128
|
+
max_new_tokens: int = 200
|
|
129
|
+
) -> torch.Tensor:
|
|
130
|
+
eos_token_id = self.tokenizer.eos_token_id
|
|
131
|
+
|
|
132
|
+
with torch.no_grad():
|
|
133
|
+
past_key_values = EncoderDecoderCache(
|
|
134
|
+
self_attention_cache=DynamicCache(),
|
|
135
|
+
cross_attention_cache=DynamicCache()
|
|
136
|
+
)
|
|
137
|
+
|
|
138
|
+
decoder_out = self.model.model.decoder(
|
|
139
|
+
input_ids=self.stable_prefix,
|
|
140
|
+
encoder_hidden_states=encoder_hidden_states,
|
|
141
|
+
past_key_values=past_key_values,
|
|
142
|
+
use_cache=True,
|
|
143
|
+
return_dict=True,
|
|
144
|
+
)
|
|
145
|
+
past_key_values = decoder_out.past_key_values
|
|
146
|
+
prefix_logits = self.model.lm_head(decoder_out.last_hidden_state)
|
|
147
|
+
next_token_id = torch.argmax(prefix_logits[:, -1, :], dim=-1).unsqueeze(-1)
|
|
148
|
+
|
|
149
|
+
generated_tokens = self.stable_prefix.clone()
|
|
150
|
+
|
|
151
|
+
for _ in range(max_new_tokens):
|
|
152
|
+
if next_token_id.item() == eos_token_id:
|
|
153
|
+
break
|
|
154
|
+
|
|
155
|
+
generated_tokens = torch.cat([generated_tokens, next_token_id], dim=-1)
|
|
156
|
+
|
|
157
|
+
decoder_out = self.model.model.decoder(
|
|
158
|
+
input_ids=next_token_id,
|
|
159
|
+
encoder_hidden_states=encoder_hidden_states,
|
|
160
|
+
past_key_values=past_key_values,
|
|
161
|
+
use_cache=True,
|
|
162
|
+
return_dict=True,
|
|
163
|
+
)
|
|
164
|
+
past_key_values = decoder_out.past_key_values
|
|
165
|
+
logits = self.model.lm_head(decoder_out.last_hidden_state)
|
|
166
|
+
next_token_id = torch.argmax(logits[:, -1, :], dim=-1).unsqueeze(-1)
|
|
167
|
+
|
|
168
|
+
return generated_tokens
|
|
169
|
+
|
|
170
|
+
|
|
171
|
+
|
|
172
|
+
if __name__ == '__main__':
|
|
173
|
+
# src_texts = [
|
|
174
|
+
# "Have you noticed how accurate",
|
|
175
|
+
# "LLM are now that GPU have became more powerful, ",
|
|
176
|
+
# "especially when we think at the difficulties we had in the 2010 era",
|
|
177
|
+
# "where online chatbots were not performant",
|
|
178
|
+
# "at all, and ofter doing strict rules worked better",
|
|
179
|
+
# "do you remember that era?"]
|
|
180
|
+
|
|
181
|
+
# src_texts = [
|
|
182
|
+
# "As-tu remarqué à quel point",
|
|
183
|
+
# "les LLM sont précis maintenant que les GPU sont devenus plus puissants, ",
|
|
184
|
+
# "surtout quand on pense aux difficultés qu'on avait dans les années 2010",
|
|
185
|
+
# "où les chatbots en ligne n'étaient pas performants",
|
|
186
|
+
# "du tout, et souvent faire des règles strictes fonctionnait mieux.",
|
|
187
|
+
# "Tu te souviens de cette époque, ",
|
|
188
|
+
# "c'était bien plus compliqué de travailler"]
|
|
189
|
+
|
|
190
|
+
translation_backend = TranslationBackend(source_lang='fra_Latn', target_lang="eng_Latn")
|
|
191
|
+
|
|
192
|
+
# # for i in range(len(src_texts)+1):
|
|
193
|
+
# # translation, buffer = translation_backend.translate(" ".join(src_texts[:i]))
|
|
194
|
+
# # print(translation, '|', buffer)
|
|
195
|
+
|
|
196
|
+
# tokens, text = translation_backend.simple_translation("Ceci est un test de traduction. Nous avons fait tout notre possible. Jusqu'à partir de 0")[1]
|
|
197
|
+
|
|
198
|
+
# # text = "Ceci est un test de traduction. Nous avons fait tout notre possible. Jusqu'à partir de 0, où nous avions pris la route, tout droit, et loin, loin, loin ! tu penses qu'on peut vraiment faire la longeur qu'on veut ?"
|
|
199
|
+
# print(output_text)
|
|
200
|
+
|
|
201
|
+
|
|
202
|
+
src_texts = ["Il s’arrêtait par moments devant une villa",
|
|
203
|
+
"coquettement nichée dans la verdure, il regardait",
|
|
204
|
+
"par la grille et voyait au loin des femmes",
|
|
205
|
+
"élégantes sur les balcons et les terrasses, des",
|
|
206
|
+
"enfants couraient dans les jardins. Il s’intéressait",
|
|
207
|
+
"surtout aux fleurs ; c’étaient elles qui attiraient",
|
|
208
|
+
"particulièrement ses regards. De temps en temps,"
|
|
209
|
+
"il voyait passer des cavaliers, des amazones et de"
|
|
210
|
+
"belles voitures ; il les suivait d’un œil curieux et",
|
|
211
|
+
"les oubliait avant qu’ils eussent disparu. "]
|
|
212
|
+
|
|
213
|
+
text = " ".join(src_texts)
|
|
214
|
+
|
|
215
|
+
|
|
216
|
+
target = """He would stop occasionally in front of a villa
|
|
217
|
+
nestled charmingly in the greenery, look
|
|
218
|
+
through the gate, and see elegant women
|
|
219
|
+
on balconies and terraces in the distance,
|
|
220
|
+
children running in the gardens. He was particularly interested
|
|
221
|
+
in the flowers; they were what caught his eye. From time to time,
|
|
222
|
+
he would see horsemen, horsewomen, and beautiful carriages passing by;
|
|
223
|
+
he would follow them with a curious eye and
|
|
224
|
+
forget them before they had disappeared.
|
|
225
|
+
"""
|
|
226
|
+
l_vals = []
|
|
227
|
+
for i in range(1, len(src_texts)+1):
|
|
228
|
+
truncated_text = " ".join(src_texts[:i])
|
|
229
|
+
input_tokens = translation_backend.tokenizer(truncated_text, return_tensors="pt").to(translation_backend.device)
|
|
230
|
+
encoder_outputs = translation_backend.model.get_encoder()(**input_tokens)
|
|
231
|
+
output_tokens = translation_backend.model.generate(
|
|
232
|
+
encoder_outputs=encoder_outputs,
|
|
233
|
+
forced_bos_token_id=translation_backend.bos_token_id
|
|
234
|
+
)
|
|
235
|
+
output_text = translation_backend.tokenizer.decode(output_tokens[0], skip_special_tokens=True)
|
|
236
|
+
l_vals.append(
|
|
237
|
+
{
|
|
238
|
+
"input": truncated_text,
|
|
239
|
+
"output_tokens_shape": output_tokens.shape[1],
|
|
240
|
+
"output_text": output_text,
|
|
241
|
+
}
|
|
242
|
+
)
|
|
243
|
+
# input_tokens = translation_backend.tokenizer(text, return_tensors="pt").to(translation_backend.device)
|
|
244
|
+
# encoder_outputs = translation_backend.model.get_encoder()(**input_tokens)
|
|
245
|
+
# output_tokens = translation_backend.model.generate(
|
|
246
|
+
# encoder_outputs=encoder_outputs,
|
|
247
|
+
# forced_bos_token_id=translation_backend.bos_token_id
|
|
248
|
+
# )
|
|
249
|
+
# output_text = translation_backend.tokenizer.decode(output_tokens[0], skip_special_tokens=True)
|
|
250
|
+
|
|
251
|
+
|
|
252
|
+
print('end')
|