stress_pattern2poem.charllama_1.3B

Requirements

Install char-level tokenizer:

pip install git+https://github.com/Koziev/character-tokenizer

Usage

Run the following code:

import torch
import transformers

# Install this package: pip install git+https://github.com/Koziev/character-tokenizer
import charactertokenizer


if __name__ == '__main__':
    generation_args = {'max_length': 1024,
                       'num_return_sequences': 5,
                       'do_sample': True,
                       'no_repeat_ngram_size': 10,
                       'temperature': 0.8,
                       'top_p': 0.8,
                       'top_k': None,
                       }

    user_prompt = """0 1 0 1 0 ёлочка
0 1 0 1 росла
0 1 0 1 0 стройная
0 1 0 0 была"""

    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    model_name = "ArsPoeticaNumeralis/stress_pattern2poem.charllama_1.3B"

    tokenizer = charactertokenizer.CharacterTokenizer.from_pretrained(model_name)
    model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
    model.to(device)

    prompt = '<s>' + user_prompt + chr(8)
    input_ids = tokenizer(prompt, return_tensors='pt').input_ids
    out_ids = model.generate(input_ids=input_ids.to(device),
                             eos_token_id=tokenizer.eos_token_id,
                             **generation_args).tolist()

    prompt_len = len(input_ids[0])
    for seq in out_ids:
        seq = seq[prompt_len:]
        output = tokenizer.decode(seq)
        if '</s>' in output:
            output = output[:output.find('</s>')].strip()
        poem_text = output
        print('-'*40)
        print(poem_text)

The output should be like this:

----------------------------------------
Ты помнишь, наша ёлочка
Стояла здесь, росла?
Стройна, красива, стройная,
Пушистая была.
----------------------------------------
Была когда-то ёлочка
Красива, да росла.
Теперь она не стройная,
Замучена была.
----------------------------------------
Совсем малютка ёлочка
Вот только что росла.
Была красива стройная,
Как девушка была.
----------------------------------------
Однажды только ёлочка
На свете всём росла.
Она была не стройная,
Невзрачная была.
----------------------------------------
Росла под нею ёлочка,
Она всегда росла,
Она была вся стройная,
Красивая была.