|
from typing import List, Dict |
|
from flair.models.language_model import LanguageModel |
|
|
|
|
|
class PreTrainedPipeline: |
|
def __init__(self, path=""): |
|
from huggingface_hub import hf_hub_download |
|
|
|
hf_hub_download(repo_id="dchaplinsky/flair-uk-forward", filename="best-lm.pt") |
|
|
|
self.model = LanguageModel.load_language_model("best-lm.pt") |
|
|
|
def __call__(self, inputs: str) -> List[Dict]: |
|
""" |
|
Args: |
|
inputs (:obj:`str`): |
|
a string containing some text |
|
Return: |
|
A :obj:`str` |
|
""" |
|
inputs = inputs.strip() |
|
return [{"generated_text": self.model.generate_text([inputs])[0]}] |
|
|