stefan-it's picture
pipeline: fix style
c9aa0de
raw
history blame
407 Bytes
from typing import List, Dict
from wtpsplit import WtP
class PreTrainedPipeline:
def __init__(self, path=""):
self.model = WtP(model_name_or_model="wtp-canine-s-de-hist-12l", hub_prefix="stefan-it")
def __call__(self, inputs: str) -> List[Dict]:
return [{"generated_text": "</s>\n".join(sent.strip() for sent in self.model.split(inputs.strip(), lang_code="de", style="digibok"))}]