style: adjust model.py
Browse files- server/model.py +10 -9
server/model.py
CHANGED
@@ -1,11 +1,12 @@
|
|
1 |
import sys, os
|
2 |
import ctranslate2
|
3 |
import sentencepiece as spm
|
|
|
4 |
|
5 |
|
6 |
|
7 |
-
def indexOf(
|
8 |
-
try: return
|
9 |
except: return -1
|
10 |
|
11 |
|
@@ -24,13 +25,13 @@ class SugoiTranslator:
|
|
24 |
return [sp.encode(text, out_type=str)]
|
25 |
|
26 |
|
27 |
-
def detokenizeBatch(self,
|
28 |
sp = spm.SentencePieceProcessor(self.sp_target_model)
|
29 |
-
|
30 |
-
return
|
31 |
|
32 |
|
33 |
-
def translate(self, text: str):
|
34 |
translated = self.translator.translate_batch(
|
35 |
source= self.tokenizeBatch(text),
|
36 |
num_hypotheses= 1,
|
@@ -46,10 +47,10 @@ class SugoiTranslator:
|
|
46 |
|
47 |
|
48 |
if __name__ == "__main__":
|
|
|
|
|
49 |
index = indexOf(sys.argv, "-modelDir")
|
50 |
-
if index != -1:
|
51 |
-
global modelDir
|
52 |
-
modelDir = sys.argv[index+1]
|
53 |
|
54 |
sugoiTranslator = SugoiTranslator(modelDir)
|
55 |
translated = sugoiTranslator.translate("ダンガンロンパ 希望の学園と絶望の高校生")
|
|
|
1 |
import sys, os
|
2 |
import ctranslate2
|
3 |
import sentencepiece as spm
|
4 |
+
from typing import Union
|
5 |
|
6 |
|
7 |
|
8 |
+
def indexOf(arr: list, value):
|
9 |
+
try: return arr.index(value)
|
10 |
except: return -1
|
11 |
|
12 |
|
|
|
25 |
return [sp.encode(text, out_type=str)]
|
26 |
|
27 |
|
28 |
+
def detokenizeBatch(self, token):
|
29 |
sp = spm.SentencePieceProcessor(self.sp_target_model)
|
30 |
+
text = sp.decode(token)
|
31 |
+
return text
|
32 |
|
33 |
|
34 |
+
def translate(self, text: Union[str, list]):
|
35 |
translated = self.translator.translate_batch(
|
36 |
source= self.tokenizeBatch(text),
|
37 |
num_hypotheses= 1,
|
|
|
47 |
|
48 |
|
49 |
if __name__ == "__main__":
|
50 |
+
global modelDir
|
51 |
+
modelDir = "./model"
|
52 |
index = indexOf(sys.argv, "-modelDir")
|
53 |
+
if index != -1: modelDir = sys.argv[index+1]
|
|
|
|
|
54 |
|
55 |
sugoiTranslator = SugoiTranslator(modelDir)
|
56 |
translated = sugoiTranslator.translate("ダンガンロンパ 希望の学園と絶望の高校生")
|