Added translation wrapper
Browse files- pipeline.py +6 -3
pipeline.py
CHANGED
@@ -18,11 +18,14 @@ class PreTrainedPipeline():
|
|
18 |
self.m2m100_tokenizer = transformers.AutoTokenizer.from_pretrained("facebook/m2m100_418M")
|
19 |
|
20 |
def __call__(self, inputs: str) -> List[Dict]:
|
21 |
-
|
|
|
22 |
# Run dialogpt
|
23 |
-
generated_text = self.dialogpt(
|
|
|
|
|
24 |
|
25 |
-
return [{"generated_text":
|
26 |
|
27 |
def dialogpt(self, inputs: str) -> str:
|
28 |
# Get input tokens
|
|
|
18 |
self.m2m100_tokenizer = transformers.AutoTokenizer.from_pretrained("facebook/m2m100_418M")
|
19 |
|
20 |
def __call__(self, inputs: str) -> List[Dict]:
|
21 |
+
# to eng
|
22 |
+
en_text = self.m2m100(inputs, "uk", "en")
|
23 |
# Run dialogpt
|
24 |
+
generated_text = self.dialogpt(en_text)
|
25 |
+
# to ukr
|
26 |
+
uk_text = self.m2m100(generated_text, "en", "uk")
|
27 |
|
28 |
+
return [{"generated_text": uk_text}]
|
29 |
|
30 |
def dialogpt(self, inputs: str) -> str:
|
31 |
# Get input tokens
|