| from transformers import AutoTokenizer, TFGPT2LMHeadModel, pipeline |
| from transformers.pipelines import TextGenerationPipeline |
| from typing import Union |
|
|
| class QuoteGenerator(): |
| def __init__(self, model_name:str='gruhit-patel/quote-generator-v2'): |
| self.model_name = model_name |
| self.quote_generator: TextGenerationPipeline |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) |
| self.model = TFGPT2LMHeadModel.from_pretrained(self.model_name) |
| self.default_tags = 'love,life' |
| print("Model has been loaded") |
| |
| def load_generator(self) -> None: |
| self.quote_generator = pipeline('text-generation', model=self.model, tokenizer=self.tokenizer) |
| print("Pipeline has been generated") |
| |
| def preprocess_tags(self, tags: Union[None, str] = None) -> str: |
| if tags is None: |
| tags = self.default_tags |
| |
| return self.tokenizer.bos_token + tags + '<bot>:' |
|
|
| def generate_quote(self, tags:Union[None, str], max_new_tokens: int, do_sample: bool, |
| num_beams: int, top_k: int, top_p: float, temperature: float): |
| |
| tags = self.preprocess_tags(tags) |
| output = self.quote_generator( |
| tags, |
| max_new_tokens=max_new_tokens, |
| num_beams=num_beams, |
| temperature=temperature, |
| top_k=top_k, |
| top_p=top_p, |
| do_sample = do_sample, |
| early_stopping=num_beams > 1 |
| ) |
| |
| return output[0]['generated_text'] |
|
|
|
|
