from transformers import pipeline, AutoTokenizer, AutoModelWithLMHead
from transformers.pipelines import TextGenerationPipeline

class QuoteGenerator:
    def __init__(self):
        self.quote_generator: TextGenerationPipeline
        self.tokenizer = AutoTokenizer.from_pretrained("nandinib1999/quote-generator")
        self.model = AutoModelWithLMHead.from_pretrained("nandinib1999/quote-generator")   
        self.default_prompts = ['life is not fair', 'she is a warrior', 'friendship is like', 'you cannot find light']

    def load_generator(self) -> None:
        self.quote_generator = pipeline('text-generation', model = self.model, tokenizer = self.tokenizer)
    def clean_text(self, quote):
        quote = quote.strip()
        # Remove any special characters
        quote = re.sub('[~`^*&%$@#|]', ' ', quote)
        # Remove extra spaces
        quote = re.sub("\s\s+", " ", quote)
        # Remove space before punctuation
        quote = re.sub('\s[.,;!?]', '\1', quote)
        # Sentence case
        quote = quote.capitalize()
        return quote

    def generate_quote(self, prompt_start: str, min_length=3, max_length=30, top_p=0.9):
        if len(prompt_start.strip()) == 0:
            rand_index = random.randint(0, len(self.default_prompts)-1)
            prompt_start = self.default_prompts[rand_index]
            
        quote_gen = self.quote_generator(prompt_start.strip(), min_length = min_length, max_length = max_length, temperature = 1.0, top_k = 5, top_p = top_p)
        quote = quote_gen[0]['generated_text']
        quote_clean = self.clean_text(quote)
        return quote_clean