Spaces:
Paused
Paused
File size: 806 Bytes
f95071d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
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
|