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']