image2poetry Pushkin Tyutchev Mayakovsky Esenin Blok

This repo contains model for russian poetry generation from images. Poetry can be generated in style of poets: Маяковский, Пушкин, Есенин, Тютчев, Блок. The model is fune-tuned concatecation of pre-trained model tuman/vit-rugpt2-image-captioning. To use this model you can write:

from PIL import Image
import requests
from transformers import AutoTokenizer, VisionEncoderDecoderModel, ViTImageProcessor

def generate_poetry(fine_tuned_model, image, tokenizer, author):
    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
    pixel_values = pixel_values.to(device)
    
    # Encode author's name and prepare as input to the decoder
    author_input = f"<bos> {author} <sep>"
    decoder_input_ids = tokenizer.encode(author_input, return_tensors="pt").to(device)

    # Generate the poetry with the fine-tuned VisionEncoderDecoder model
    generated_tokens = fine_tuned_model.generate(
        pixel_values,
        decoder_input_ids=decoder_input_ids,
        max_length=300,
        num_beams=3,
        top_p=0.8,
        temperature=2.0,
        do_sample=True,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )

    # Decode the generated tokens
    generated_poetry = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
    generated_poetry = generated_poetry.split(f'{author}')[-1]
    return generated_poetry

path = 'AnyaSchen/vit-rugpt3-large-poetry-ft'
fine_tuned_model = VisionEncoderDecoderModel.from_pretrained(path).to(device)
feature_extractor = ViTImageProcessor.from_pretrained(path)
tokenizer = AutoTokenizer.from_pretrained(path)


url = 'https://anandaindia.org/wp-content/uploads/2018/12/happy-man.jpg'
image = Image.open(requests.get(url, stream=True).raw)

generated_poetry = generate_poetry(fine_tuned_model, image, tokenizer, 'Маяковский')
print(generated_poetry)