|
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 |
|
|
|
self.model = LanguageModel.load_language_model( |
|
hf_hub_download(repo_id="dchaplinsky/flair-uk-backward-large", filename="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, temperature=0.5)[0]}] |
|
|