theodotus commited on
Commit
319ae94
1 Parent(s): fcae46c

Added translation wrapper

Browse files
Files changed (1) hide show
  1. 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(inputs)
 
 
24
 
25
- return [{"generated_text": 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