import transformers
from transformers import (
        AutoModelForSeq2SeqLM,
        Seq2SeqTrainingArguments,
        Seq2SeqTrainer,
        AutoTokenizer
)
from transformers import pipeline


class ModelAPI:
    def __init__(self):
        path = 't5_samsum/checkpoint-9210'
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(path)
        self.model.to("cpu")
        self.summarizer_pipeline = pipeline("summarization",
                                model = self.model,
                                tokenizer = self.tokenizer,
                                device = "cpu")

    def create_summary(self, input_text, model_pipeline):
        summary = model_pipeline(input_text)
        return summary

    def get_summary(self, text):
        summary = self.create_summary(text, self.summarizer_pipeline)
        return summary[0]