from transformers import AutoTokenizer, AutoModelForCausalLM


class LyricGeneratorModel:
    def __init__(self, repo_id: str):
        self.model = AutoModelForCausalLM.from_pretrained(repo_id)

        self.tokenizer = AutoTokenizer.from_pretrained(repo_id)
        self.tokenizer.truncation_side = "right"
        self.tokenizer.pad_token = self.tokenizer.eos_token
        self.tokenizer.padding_side = "right"

    def generate_lyrics(self, prompt: str, max_length: int):
        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
        input_ids = input_ids.to("cuda")

        output_tokens = self.model.generate(
            input_ids, do_sample=True, max_length=max_length
        )

        output_text = self.tokenizer.batch_decode(output_tokens)[0]

        return output_text