Описание модели

RUDOLPH — это авторегрессионная модель для генерации текста и изображения. Генеративная архитектура модели, а также особенность деления входной последовательности на сегменты токенов, отвечающих за кодирование текста и изображения, позволяют модели решать различные комбинации задач в доступных модальностях: от ответа на вопросы по тексту и изображению до генерации изображений по текстовым запросам.

Схема модели представлена на рисунке ниже.

Модель RUDOLPH 2.7B работает с последовательностью токенов, которую можно разделить на 3 основных сегмента (384 левых текстовых токена), 576 токенов для кодирования изображения, 128 правых текстовых токенов. Для токенизации изображения используется Sber-VQ-GAN, токенизация текстовой последовательностей осуществляется с помощью токенизатора YTTM.

Для того чтобы модель могла эффективно обрабатывать токены, относящиеся к двум модальностям, в архитектуре RUDOLPH реализован многоуровневый механизм внимания: 1) маскированный механизм внимания (masked self-attention) для текстовых токенов; 2) механизм разреженного внимания (sparse attention) для визуальных токенов. Ефким образом, когда генерируется новый токен изображения открытыми остаются не все сгенерированные до этого токены, а лишь его соседи в строке, столбце и некотором контекстом окне.

Маски внимания, реализованные в модели RUDOLPH 2.7B приведены на рисунке.

Модель была обучена на 119M пар текст-изображение и 60M текстов на русском языке.

В процессе предобучения модель видела три типах задач: text2text (Language Modeling в левых текстовых токенах), image2text (image captioning), text2image (image generation). Таким образом, предобученная модель RUDOLPH умеет создавать описание к изображению, продолжать текст по затравке, а также генерировать изображение по текстовому запросу.

Предобученную модель можно использовать в zero-shot режиме, либо дообучить на необходимых дополнительных задачах для достижения лучшего результата. Пример дообучения модели на задачу Machine Reading Comprehension можно найти в github RUDOLPH.

Преимущества

Основным преимуществом модели является ее универсальность с точки зрения решения разных комбинаций задач в бимодальном режиме. В таблице ниже приведены примеры таких задач. Строки отвечают за модальность входа, колонки за модальность выхода.

В таблице мы привели наиболее яркие примеры решаемых задач, вы можете добавить свои задачи, например, задачи в более узком домене.

Также отметим, что бинаризация входа и выхода является упрощением, на самом деле, RUDOLPH умеет решать задачи с бимодальным входом (text-image→text), а также с бимодальным выходом (например, text→image+text).

Наличие текстовых токенов в начале входной последовательности позволяет нам задавать тип решаемой задачи на естественном языке (в виде инструкции), например, “Распознай текст на изображении.”, “Сгенерируй изображение по тексту.”, “Ответь на вопрос по тексту.” и т.д.

Сценарии использования

Посмотрим на примеры генерацию RUDOLPH 2.7B после предобучения.

Генерация изображения по запросу Морской пейзаж в закате солнца.

Math QA (few-shot режим) **|3+4=7|2+1=3|4+3=7|8+1=9|3+2=5|2+2=4|1+6=7|3+2=5|9+2=**11|

Генерация текста

из окна открывался вид* на морское побережье и горы*.

Image captioning

на картинке изображены три белые и черные птицы.

Информация об использовании модели

KFServing

Класс KFServingRUDOLPHModel представлен ниже.

Вы можете подавать на вход модели:

• тип задачи: “T2T” (text generation), “T2I” (image generation), “I2T” (image captioning)
• текст: текстовое описание задачи
• картинки в формате base64

На выходе модель выдаст текст, соответствующий входным данным, или изображение, соответствующее текстовому запросу.

import numpy as np
from PIL import Image
import os
import cv2
import re
import base64
from io import BytesIO
import time

import kfserving

import torch
import ruclip

import sys

from rudalle.image_prompts import ImagePrompts
from rudolph.model import get_rudolph_model, ruDolphModel, FP16Module
from rudolph import utils
from rudolph.model.utils import get_attention_mask
from rudalle import get_tokenizer
from vae import get_vae
from rudolph.api import ruDolphApi

import requests
from PIL import Image, ImageDraw
import io
from app.utils import *
    
class KFServingRUDOLPHModel(kfserving.KFModel):

    def __init__(self, name: str):
        super().__init__(name)
        self.name = name

        self.device = torch.device('cuda:0')
        self.cache_dir = './weights'
        self.vae_dir = './weights/vqgan.gumbelf8-sber.model.ckpt'
        self.tokenizer_weights = './weights/bpe.model'
        
        self.gpu = True
        self.ready = False

    def load(self):

        self.model_fbc2 = get_rudolph_model('2.7B', pretrained=False, fp16=True, device=self.device)
        self.model_fbc2.load_state_dict(torch.load(f'{self.cache_dir}/pytorch_model.bin'))

        self.tokenizer = get_tokenizer(path=self.tokenizer_weights)
        self.vae = get_vae(dwt=False, filename=self.vae_dir).to(self.device)

        self.api = ruDolphApi(self.model_fbc2, self.tokenizer, self.vae, bs=48)
        
        self.ready = True

    def image_generation(self, text):
        images_num, codebooks_num = 1, 12
        top_k, top_p = 1024, 0.99
        pil_images, scores = self.api.dream(
            text=text, images_num=images_num, codebooks_num=codebooks_num,
            top_k=top_k, top_p=top_p, special_token='<LT_T2I>'
        )

        return {"ok": True, "image": encode_img(np.array(pil_images[0]))}

    def captioning(self, pil_img):
        result = self.api.image_captioning(
            pil_img, early_stop=32, generations_num=48, captions_num=5, seed=42,
            l_special_token='<LT_I2T>'
        )
        return {"ok": True, "text": result[0]["text"]}

    def text_generation(self, text):
        top_k, top_p = 64, 0.8
        texts_num = 48
        temperature = 0.9

        text_answer = self.api.generate_texts(
            template=text,
            top_k=top_k, top_p=top_p, texts_num=texts_num, seed=42, temperature=temperature,
            special_token='<LT_T2T>',
            early_stop=32,
        )[:5]
        return {"ok": True, "text": text_answer[0]}
    
    def predict(self, request):
        request = request['instances'][0]
        try:
            text = request['text']
            image = None
            if 'image' in request and request['image'] is not None:
                image = decode_img(request['image'])

            if request['task'] == 'T2I':
                result_data = self.image_generation(text)
            elif request['task'] == 'I2T':
                result_data = self.captioning(image)
            elif request['task'] == 'T2T':
                result_data = self.text_generation(text)
            else:
                return {"ok": False, "error": "Unknown task"}
        except Exception as err:
            info_str = f"query={request['text']} | image={image is not None}"
            return {"ok": False, "error": f'Error while generating: {err}
 INFO: {info_str}'}
        return result_data

Функция predict возвращает словарь с полями, соответствующим типам задач (request[”task”]):

{'ok': True,
 'text': 'Вид на морское побережье и горы из окна'})

Пример работы с моделью

from rudolph_deploy import KFServingRUDOLPHModel

model = KFServingRUDOLPHModel("kfserving-RUDOLPH")
model.load()

# zero-shot text generation
result = model.predict({
     "instances": [
     {
     "task": "T2T",
     "text": 'Вид на морское побережье и '
     }
     ]
     }
)

"Res: ", result

Ccылки

Github HuggingFace 🤗