Spaces:
Paused
Paused
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 | |