Almost almost there yet yet

pipeline.py

@@ -10,7 +10,7 @@ class PreTrainedPipeline:
         # This function is only called once, so do all the heavy processing I/O here"""
         self.model = PunctuationCapitalizationModel.from_pretrained("dchaplinsky/punctuation_uk_bert")
 
-    def __call__(self, inputs: str) -> List[str]:
+    def __call__(self, inputs: str) -> Dict[str, str]:
         """
         Args:
             inputs (:obj:`str`):
@@ -19,4 +19,4 @@ class PreTrainedPipeline:
             A :obj:`str`
         """
         inputs = inputs.strip()
-        return {"generated_text": [self.model.add_punctuation_capitalization(inputs)]}
+        return {"generated_text": self.model.add_punctuation_capitalization([inputs])[0]}