demo init

.gitignore

@@ -0,0 +1,3 @@
+data
+flagged
+__pycache__

app.py

@@ -0,0 +1,43 @@
+from pathlib import Path
+
+import gdown
+# from PIL import Image
+# from numpy import asarray
+
+from run import launch
+
+# download
+if not Path('./data').is_dir():
+    url = 'https://drive.google.com/drive/folders/1hfHWDn5iXsdjB63E5zdZBAoRLWHQC3LD'
+    gdown.download_folder(url, quiet=True, use_cookies=False, output="./data/")
+
+# example image from COCO data
+image_urls = {
+    '108953': 'https://farm8.staticflickr.com/7160/6484651991_9d1eaa557a_z.jpg'
+}
+images = {}
+for k, url in image_urls.items():
+    ext = Path(url).suffix
+    output = Path(f"data/images/{k}{ext}")
+    if not output.is_file():
+        output.parent.mkdir(exist_ok=True)
+        gdown.download(url, quiet=True, use_cookies=False, output=str(output))
+    images[k] = str(output)
+
+'''
+for k, v in images.items():
+    with Image.open(v) as image:
+        # image = asarray(image)
+        images[k] = image
+'''
+
+
+examples = [[
+    v,
+    'My favourite recipe:',
+    20,
+    10,
+    False
+] for v in images.values()]
+
+launch(examples)

arguments.py

@@ -0,0 +1,58 @@
+import os
+import json
+import argparse
+import logging
+from pathlib import Path
+
+import torch
+
+
+logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))
+log = logging.getLogger(__name__)
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='ESPER')
+
+
+    parser.add_argument(
+        '--init-model', type=str, default='gpt2', help='language model used for policy.')
+    parser.add_argument(
+        '--label_path', type=str, default='./data/esper_demo/labels_all.json', help='style label info file path')
+    parser.add_argument(
+        '--checkpoint', type=str, default='./data/esper_demo/ckpt/gpt2_style', help='checkpoint file path')
+
+    parser.add_argument(
+        '--prefix_length', type=int, default=10, help='prefix length for the visual mapper')
+    parser.add_argument(
+        '--clipcap_num_layers', type=int, default=1, help='num_layers for the visual mapper')
+    parser.add_argument(
+        '--use_transformer_mapper', action='store_true', default=False, help='use transformer mapper instead of mlp')
+    parser.add_argument(
+        '--use_label_prefix', action='store_true', default=False, help='label as prefixes')
+    parser.add_argument(
+        '--clip_model_type', type=str, default='ViT-B/32', help='clip backbone type')
+
+    parser.add_argument(
+        '--infer_no_repeat_size', type=int, default=2, help="no repeat ngram size for inference")
+    parser.add_argument(
+        '--response-length', type=int, default=20, help='number of tokens to generate for each prompt.')
+    parser.add_argument(
+        '--num-gpus', type=int, default=None, help='number of gpus. use all available if none')
+    parser.add_argument(
+        '--port', type=int, default=None, help="port for the demo server")
+
+    args = parser.parse_args()
+    args.cuda = torch.cuda.is_available()
+
+    if args.use_label_prefix:
+        log.info(f'using label prefix')
+    num_gpus = torch.cuda.device_count()
+    if args.num_gpus is None:
+        args.num_gpus = num_gpus
+    else:
+        args.num_gpus = min(num_gpus, args.num_gpus)
+
+    if args.checkpoint is not None:
+        args.checkpoint = str(Path(args.checkpoint).resolve())
+    return args

clipcap.py

@@ -0,0 +1,385 @@
+import os
+import math
+import logging
+import json
+from pathlib import Path
+from typing import Tuple, Optional, Union
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from transformers import AutoModelForCausalLM, AutoTokenizer, GPTJForCausalLM
+
+
+logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))
+log = logging.getLogger(__name__)
+
+
+def load_weights(self, Module, path, name, default_name, prev_name=None, **kwargs):
+    hparams = None
+    assert isinstance(default_name, str), f'invalid default transformer name: {default_name}'
+    model = get_transformer_module(Module, default_name, **kwargs)
+    setattr(self, name, model)
+    return hparams
+
+
+def get_transformer_module(Module, default_name, **kwargs):
+    if default_name == 'EleutherAI/gpt-j-6B':
+        kwargs = {**kwargs, **dict(revision="float16", torch_dtype=torch.float16, low_cpu_mem_usage=True)}
+    model = Module.from_pretrained(default_name, **kwargs)
+    return model
+
+
+class MLP(nn.Module):
+    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
+        super(MLP, self).__init__()
+
+        self.divider = math.sqrt(sizes[-1] / sizes[0])
+        layers = []
+        for i in range(len(sizes) - 1):
+            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
+            if i < len(sizes) - 2:
+                layers.append(act())
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x / self.divider  # scaling for the initial stability
+        x = self.model(x)
+        return x
+
+
+class MlpTransformer(nn.Module):
+    def __init__(self, in_dim, h_dim, out_d: Optional[int] = None, act=F.relu, dropout=0.):
+        super().__init__()
+        out_d = out_d if out_d is not None else in_dim
+        self.fc1 = nn.Linear(in_dim, h_dim)
+        self.act = act
+        self.fc2 = nn.Linear(h_dim, out_d)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, dim_self, dim_ref, num_heads, bias=True, dropout=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim_self // num_heads
+        self.scale = head_dim ** -0.5
+        self.to_queries = nn.Linear(dim_self, dim_self, bias=bias)
+        self.to_keys_values = nn.Linear(dim_ref, dim_self * 2, bias=bias)
+        self.project = nn.Linear(dim_self, dim_self)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, y=None, mask=None):
+        y = y if y is not None else x
+        b, n, c = x.shape
+        _, m, d = y.shape
+        # b n h dh
+        queries = self.to_queries(x).reshape(b, n, self.num_heads, c // self.num_heads)
+        # b m 2 h dh
+        keys_values = self.to_keys_values(y).reshape(b, m, 2, self.num_heads, c // self.num_heads)
+        keys, values = keys_values[:, :, 0], keys_values[:, :, 1]
+        attention = torch.einsum('bnhd,bmhd->bnmh', queries, keys) * self.scale
+        if mask is not None:
+            if mask.dim() == 2:
+                mask = mask.unsqueeze(1)
+            attention = attention.masked_fill(mask.unsqueeze(3), float("-inf"))
+        attention = attention.softmax(dim=2)
+        out = torch.einsum('bnmh,bmhd->bnhd', attention, values).reshape(b, n, c)
+        out = self.project(out)
+        return out, attention
+
+
+class TransformerLayer(nn.Module):
+    def __init__(self, dim_self, dim_ref, num_heads, mlp_ratio=4., bias=False, dropout=0., act=F.relu,
+                 norm_layer: nn.Module = nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim_self)
+        self.attn = MultiHeadAttention(dim_self, dim_ref, num_heads, bias=bias, dropout=dropout)
+        self.norm2 = norm_layer(dim_self)
+        self.mlp = MlpTransformer(dim_self, int(dim_self * mlp_ratio), act=act, dropout=dropout)
+
+    def forward_with_attention(self, x, y=None, mask=None):
+        x_, attention = self.attn(self.norm1(x), y, mask)
+        x = x + x_
+        x = x + self.mlp(self.norm2(x))
+        return x, attention
+
+    def forward(self, x, y=None, mask=None):
+        x = x + self.attn(self.norm1(x), y, mask)[0]
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def forward_with_attention(self, x, y=None, mask=None):
+        attentions = []
+        for layer in self.layers:
+            x, att = layer.forward_with_attention(x, y, mask)
+            attentions.append(att)
+        return x, attentions
+
+    def forward(self, x, y=None, mask=None):
+        for i, layer in enumerate(self.layers):
+            if i % 2 == 0 and self.enc_dec: # cross
+                x = layer(x, y)
+            elif self.enc_dec:  # self
+                x = layer(x, x, mask)
+            else:  # self or cross
+                x = layer(x, y, mask)
+        return x
+
+    def __init__(self, dim_self: int, num_heads: int, num_layers: int, dim_ref: Optional[int] = None,
+                 mlp_ratio: float = 2., act=F.relu, norm_layer: nn.Module = nn.LayerNorm, enc_dec: bool = False):
+        super(Transformer, self).__init__()
+        dim_ref = dim_ref if dim_ref is not None else dim_self
+        self.enc_dec = enc_dec
+        if enc_dec:
+            num_layers = num_layers * 2
+        layers = []
+        for i in range(num_layers):
+            if i % 2 == 0 and enc_dec:  # cross
+                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
+            elif enc_dec:  # self
+                layers.append(TransformerLayer(dim_self, dim_self, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
+            else:  # self or cross
+                layers.append(TransformerLayer(dim_self, dim_ref, num_heads, mlp_ratio, act=act, norm_layer=norm_layer))
+        self.layers = nn.ModuleList(layers)
+
+
+class TransformerMapper(nn.Module):
+    def __init__(self, dim_clip: int, dim_embedding: int, prefix_length: int = 10,
+                 clip_length: int = 10, num_layers: int = 8):
+        super(TransformerMapper, self).__init__()
+        self.clip_length = clip_length
+        self.transformer = Transformer(dim_embedding, 8, num_layers)
+        self.linear = nn.Linear(dim_clip, clip_length * dim_embedding)
+        self.prefix_const = nn.Parameter(torch.randn(prefix_length, dim_embedding), requires_grad=True)
+
+    def forward(self, x):
+        x = self.linear(x).view(x.shape[0], self.clip_length, -1)
+        prefix = self.prefix_const.unsqueeze(0).expand(x.shape[0], *self.prefix_const.shape)
+        prefix = torch.cat((x, prefix), dim=1)
+        out = self.transformer(prefix)[:, self.clip_length:]
+        return out
+
+
+class ClipCap(nn.Module):
+    def __init__(self, model_name, device, prefix_length: int = 10, clip_length: int = 40, prefix_size: int = 512,
+                 num_layers: int = 1, model_path: str = '', fix_gpt: bool = False,
+                 use_label_prefix: bool = False, label_path: str = '', label_length: int = 10,
+                 use_transformer_mapper: bool = False, use_ptuning_v2: bool = False,
+                 dropout: float = 0,
+                 model_weight: str = '', scalar_output: bool = False):
+        super(ClipCap, self).__init__()
+
+        self.prefix_length = prefix_length
+        self.prefix_size = prefix_size
+        self.label_length = label_length
+        self.scalar_output = scalar_output
+        self.num_layers = num_layers
+        self.use_transformer_mapper = use_transformer_mapper
+        self.use_ptuning_v2 = use_ptuning_v2
+
+        self.dropout = nn.Dropout(dropout)
+
+        hparams = load_weights(self, AutoModelForCausalLM, model_weight, 'gpt', model_name,
+                               prev_name='model')
+
+        self.device = device
+        self.gpt = self.gpt.to(self.device)
+
+        config = self.gpt.config
+        self.match_n_layer = getattr(config, 'n_layer', getattr(config, 'num_layers', None))  # gpt2 vs. gpt_neo
+        self.match_n_head = getattr(config, 'n_head', getattr(config, 'num_heads', None))
+        self.n_embd = getattr(config, 'n_embd', getattr(config, 'hidden_size', None))
+        self.match_n_embd = self.n_embd // self.match_n_head
+
+        self.clip_project = self.get_mapper()
+
+        if Path(label_path).is_file():
+            with open(label_path) as f:
+                labels = json.load(f)
+            self.labels = {i: v for v, i in labels.items()}
+            if not use_label_prefix:
+                log.info("adding label projections")
+                self.label_project = nn.Sequential(
+                    nn.Embedding(len(self.labels), self.prefix_size),
+                    self.get_mapper()
+                )
+
+        if os.path.isfile(model_path):
+            log.info(f"loading model from {model_path}")
+            weight = torch.load(model_path, map_location=torch.device('cpu'))
+            weight = {k[len('clip_project.'):]: v for k, v in weight.items()
+                      if k.startswith('clip_project.')}
+            self.clip_project.load_state_dict(weight)
+
+        if fix_gpt:
+            log.info("fixing gpt parameters")
+            for param in self.gpt.parameters():
+                param.requires_grad_(False)
+
+        if self.scalar_output:
+            self.gpt.lm_head = nn.Linear(self.gpt.transformer.embed_dim, 1).to(self.device)
+
+        self.clip_project = self.clip_project.to(self.device)
+        if hasattr(self, 'label_project'):
+            self.label_project = self.label_project.to(self.device)
+
+    def get_mapper(self):
+        if self.use_ptuning_v2:
+            total_embd = self.match_n_layer * 2 * self.n_embd
+            module = MLP((self.prefix_size,
+                          *[self.prefix_size
+                            for i in range(self.num_layers)],
+                          total_embd * self.prefix_length))
+        elif self.use_transformer_mapper:
+            log.info("using transformer mapper")
+            module = TransformerMapper(self.prefix_size, self.n_embd,
+                                       self.prefix_length, self.prefix_length, num_layers=self.num_layers)  # 8)
+        else:
+            module = MLP((self.prefix_size,
+                          *[(self.n_embd * self.prefix_length) // 2
+                            for i in range(self.num_layers)],
+                          self.n_embd * self.prefix_length))
+        return module
+
+    def get_encoder_loss(self, input_ids: torch.Tensor, features: torch.Tensor,
+                         device = None):
+        input_ids = input_ids[:, :self.prefix_length].to(device)
+        embedding = self.gpt.transformer.wte(input_ids)
+        features = features.to(device)
+        prefix_projections = self.clip_project(features.type_as(embedding)).reshape(-1, self.prefix_length, self.n_embd)
+        fct = nn.MSELoss()
+        loss = fct(prefix_projections, embedding.detach())
+        return loss
+
+    def forward(self, *args, **kwargs):
+        if self.use_ptuning_v2:
+            return self.forward_prefix(*args, **kwargs)
+        else:
+            return self.forward_embedding(*args, **kwargs)
+
+    def forward_embedding(self, input_ids: torch.Tensor, features: torch.Tensor,
+                attention_mask: Optional[torch.Tensor] = None,
+                labels: Optional[torch.Tensor] = None,
+                past_key_values = None, device = None, **kwargs):
+
+        if device is None:
+            device = self.device
+        input_ids = input_ids.to(device)
+        if features is not None:
+            features = features.to(device)
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(device)
+        if labels is not None:
+            labels = labels.to(device)
+        use_labels = labels is not None and hasattr(self, 'label_project')
+
+        embedding = self.gpt.transformer.wte(input_ids)
+        embed_txt = embedding
+        prefix_length = self.prefix_length
+        if use_labels:
+            prefix_length += self.label_length
+        if past_key_values is None:
+            prefix_projections = self.clip_project(features.type_as(embedding)).reshape(-1, self.prefix_length, self.n_embd)
+            if use_labels:
+                label_projections = self.label_project(labels.long()).reshape(-1, self.label_length, self.n_embd)
+                prefix_projections = torch.cat((prefix_projections, label_projections), dim=1)
+            embedding = torch.cat((prefix_projections.to(embedding.dtype), embedding), dim=1)
+        if torch.is_tensor(attention_mask):
+            prefix_mask = torch.ones_like(attention_mask)[:, :1].repeat(1, prefix_length)
+            attention_mask = torch.cat([prefix_mask, attention_mask], dim=1)
+        outputs = self.gpt(inputs_embeds=embedding, attention_mask=attention_mask,
+                           past_key_values=past_key_values,
+                           return_dict=True,
+                           output_attentions=False,
+                           output_hidden_states=True)
+        if past_key_values is None:
+            outputs.logits = outputs.logits[:, prefix_length:]
+        return outputs
+
+    def forward_prefix(self, input_ids: torch.Tensor, features: torch.Tensor,
+                attention_mask: Optional[torch.Tensor] = None,
+                labels: Optional[torch.Tensor] = None,
+                past_key_values = None, device = None, **kwargs):
+
+        if device is None:
+            device = self.device
+        input_ids = input_ids.to(device)
+        if features is not None:
+            features = features.to(device)
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(device)
+        if labels is not None:
+            labels = labels.to(device)
+        use_labels = labels is not None and hasattr(self, 'label_project')
+
+        prefix_length = self.prefix_length
+        if use_labels:
+            prefix_length += self.label_length
+        if past_key_values is None:
+            prefix_projections = self.clip_project(features.type_as(self.clip_project.model[0].weight))
+            prefix_projections = prefix_projections.reshape(-1, self.prefix_length,
+                                                            self.match_n_layer * 2, self.match_n_head, self.match_n_embd)
+            if use_labels:
+                label_projections = self.label_project(labels.long())
+                label_projections = label_projections.reshape(-1, self.label_length,
+                                                              self.match_n_layer * 2, self.match_n_head, self.match_n_embd)
+                prefix_projections = torch.cat((prefix_projections, label_projections), dim=1)
+            temp_control = prefix_projections
+            temp_control = self.dropout(temp_control)
+            past_key_values = temp_control.permute([2, 0, 3, 1, 4]).split(2)
+
+        if torch.is_tensor(attention_mask):
+            prefix_mask = torch.ones_like(attention_mask)[:, :1].repeat(1, prefix_length)
+            attention_mask = torch.cat([prefix_mask, attention_mask], dim=1)
+        outputs = self.gpt(input_ids=input_ids, attention_mask=attention_mask,
+                           past_key_values=past_key_values,
+                           return_dict=True,
+                           output_attentions=False,
+                           output_hidden_states=True)
+        if past_key_values is None:
+            outputs.logits = outputs.logits[:, prefix_length:]
+        return outputs
+
+    def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
+        token_type_ids = kwargs.get("token_type_ids", None)
+        # only last token for inputs_ids if past is defined in kwargs
+        if past:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+        attention_mask = kwargs.get("attention_mask", None)
+        position_ids = kwargs.get("position_ids", None)
+        features = kwargs.get("features", None)
+        labels = kwargs.get("labels", None)
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+        else:
+            position_ids = None
+
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past,
+            "use_cache": kwargs.get("use_cache"),
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "token_type_ids": token_type_ids,
+            "features": features,
+            "labels": labels,
+        }

load.py

@@ -0,0 +1,60 @@
+import os
+import logging
+import json
+from pathlib import Path
+
+import yaml
+import torch
+
+from policy import Policy
+
+
+logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))
+log = logging.getLogger(__name__)
+
+
+def load_model_args(args):
+    checkpoint = Path(args.checkpoint + '.ckpt')
+    assert checkpoint.is_file(), f"no checkpoint file: {args.checkpoint}"
+    args_path = Path(args.checkpoint + '.json')
+    if args_path.is_file():
+        with open(args_path) as f:
+            hparams = json.load(f)
+    else:
+        args_path = Path(args.checkpoint + '.yaml')
+        with open(args_path) as f:
+            hparams = yaml.safe_load(f)
+    for key in ['init_model', 'clip_model_type', 'use_caption', 'use_style_reward', 'use_transformer_mapper',
+                'prefix_length', 'clipcap_num_layers', 'use_ptuning_v2']:
+        if key in hparams:
+            setattr(args, key, hparams[key])
+    args.loaded_init_model = True
+    return args
+
+
+def load_model(args, device, finetune=False):
+    log.info('loading model')
+    policy = Policy(model_name=args.init_model, temperature=1.0, device=device,
+                    clipcap_path='None', fix_gpt=True,
+                    label_path=args.label_path,
+                    prefix_length=args.prefix_length,
+                    clipcap_num_layers=args.clipcap_num_layers,
+                    use_transformer_mapper=args.use_transformer_mapper,
+                    model_weight='None', use_label_prefix=args.use_label_prefix)
+    ckpt = args.checkpoint + '.ckpt'
+    state = torch.load(ckpt)
+    policy_key = 'policy_model'
+    if policy_key in state:
+        policy.model.load_state_dict(state[policy_key])
+    else:
+        weights = state['state_dict']
+        key = 'policy.model.'
+        if not any(k for k in weights.keys() if k.startswith(key)):
+            key = 'model.model.'
+        weights = {k[len(key):]: v for k, v in weights.items() if k.startswith(key)}
+        # weights = {k: v for k, v in weights.items() if k.startswith('clip_project.')}
+        policy.model.load_state_dict(weights, strict=False)
+    model = policy
+
+    model = model.to(device)
+    return model

policy.py

@@ -0,0 +1,219 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from typing import Union, List, Dict, Optional
+from transformers import AutoModelForCausalLM, AutoTokenizer, GPTJForCausalLM
+from transformers.generation_logits_process import (
+    LogitsProcessorList,
+    NoBadWordsLogitsProcessor,
+    NoRepeatNGramLogitsProcessor,
+)
+
+from utils import (
+    NEGATIVE_INF, HALF_NEGATIVE_INF,
+    logits_to_entropy, mask_pad
+)
+from clipcap import ClipCap
+
+
+class Policy(nn.Module):
+    def __init__(self, model_name, temperature, device, clipcap_path='', fix_gpt=False,
+                 use_transformer_mapper: bool = False, use_ptuning_v2: bool = False,
+                 prefix_length=10, clipcap_num_layers: int = 1,
+                 label_path: str = '', model_weight: str = 'None', use_label_prefix: bool = False):
+        super().__init__()
+
+        self.device = device
+
+        self.model = ClipCap(model_name, device,
+                             model_path=clipcap_path, fix_gpt=fix_gpt,
+                             prefix_length=prefix_length,
+                             num_layers=clipcap_num_layers,
+                             label_path=label_path, model_weight=model_weight,
+                             use_transformer_mapper=use_transformer_mapper,
+                             use_ptuning_v2=use_ptuning_v2,
+                             use_label_prefix=use_label_prefix)
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, pad_token="<|endoftext|>")
+        self.model.gpt.config.pad_token_id = self.tokenizer.pad_token_id
+
+        self.temperature = temperature
+
+    def get_processor(self, no_repeat_ngram_size: int = 3):
+        logits_processor = LogitsProcessorList()
+        if no_repeat_ngram_size > 0:
+            logits_processor.append(NoRepeatNGramLogitsProcessor(ngram_size=no_repeat_ngram_size))
+        '''
+        logits_processor.append(NoBadWordsLogitsProcessor([[self.tokenizer.pad_token_id]],
+                                                          self.tokenizer.pad_token_id))
+        '''
+        return logits_processor
+
+    def sample(self,
+               input_ids: torch.Tensor = None,
+               features: torch.Tensor = None,
+               attention_mask: torch.Tensor = None,
+               labels: Optional[torch.Tensor] = None,
+               max_len: int = 20,
+               sample: bool = True,
+               top_k: int = None,
+               top_p: float = None,
+               temperature: float = None,
+               no_repeat_ngram_size: int = 0,
+               invalidate_eos: bool = True,
+               device = None) -> Dict[str, Union[torch.Tensor, List[str]]]:
+        if device is None:
+            device = self.device
+        if temperature is None:
+            temperature = self.temperature
+
+        input_ids = input_ids.to(device)
+        attention_mask = attention_mask.to(device)
+
+        model_kwargs = {'attention_mask': attention_mask}
+        batch_size, input_seq_len = input_ids.shape
+
+        logits_processor = self.get_processor(no_repeat_ngram_size=no_repeat_ngram_size)
+
+        logits_warper = self.model.gpt._get_logits_warper(
+            top_k=top_k, top_p=top_p, temperature=temperature, num_beams=1
+        )
+
+        unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=device)
+        output_logprob = torch.zeros([batch_size, 0], device=device)
+        eos_logprobs = torch.zeros([batch_size, 0], device=device)
+        output_mask = torch.ones([batch_size, 0], dtype=torch.long, device=device)
+
+        self.model.eval()
+        with torch.no_grad():
+            for step in range(max_len):
+                # prepare model inputs
+                model_inputs = self.model.prepare_inputs_for_generation(input_ids,
+                                                                        features=features,
+                                                                        labels=labels,
+                                                                        **model_kwargs)
+
+                # forward pass to get next token
+                outputs = self.model(
+                    **model_inputs,
+                    device=device
+                )
+
+                # in the first decoding step, we want to use the 'real' last position for each sentence
+                if step == 0:
+                    last_non_masked_idx = torch.sum(attention_mask, dim=1) - 1
+                    next_token_logits = outputs.logits[range(batch_size), last_non_masked_idx, :]
+                else:
+                    next_token_logits = outputs.logits[:, -1, :]
+
+                negative_inf = HALF_NEGATIVE_INF if next_token_logits.dtype == torch.half else NEGATIVE_INF
+                next_token_scores = logits_processor(input_ids, next_token_logits)
+                if invalidate_eos:
+                    next_token_scores[:, self.tokenizer.eos_token_id] = negative_inf  # no endoftext
+                log_prob = F.log_softmax(next_token_scores, dim=-1)  # authentic sampling distribution
+                next_token_scores = logits_warper(input_ids, next_token_scores)
+                if sample:
+                    # Temperature (higher temperature => more likely to sample low probability tokens)
+                    probs = F.softmax(next_token_scores, dim=-1)
+                    next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+                else:
+                    # Greedy decoding
+                    next_tokens = torch.argmax(next_token_scores, dim=-1)
+
+                # finished sentences should have their next token be a padding token
+                next_tokens = next_tokens * unfinished_sequences + self.tokenizer.pad_token_id * (1 - unfinished_sequences)
+
+                    # update output mask
+                output_mask = torch.cat([output_mask, unfinished_sequences[:, None]], dim=-1)
+                # update output log probability
+                eos_logprob = log_prob[:, self.tokenizer.eos_token_id]
+                eos_logprob = eos_logprob * unfinished_sequences + negative_inf * (1 - unfinished_sequences)
+                eos_logprobs = torch.cat([eos_logprobs, eos_logprob[:, None]], dim=-1)
+
+                token_logprob = torch.gather(log_prob, 1, next_tokens[:, None]).squeeze(1)
+                token_logprob = token_logprob * unfinished_sequences + negative_inf * (1 - unfinished_sequences)
+                output_logprob = torch.cat([output_logprob, token_logprob[:, None]], dim=-1)
+
+                # update generated ids, model inputs for next step
+                input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
+                model_kwargs = self.model.gpt._update_model_kwargs_for_generation(
+                    outputs, model_kwargs, is_encoder_decoder=self.model.gpt.config.is_encoder_decoder
+                )
+
+                # if eos_token was found in one sentence, set sentence to finished
+                unfinished_sequences = unfinished_sequences.mul((next_tokens != self.tokenizer.eos_token_id).long())
+
+                if unfinished_sequences.max() == 0:
+                    break
+
+        response_ids = input_ids[:, input_seq_len:]
+        response_text = [self.tokenizer.decode(output, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+                         for output in response_ids]
+
+        prompt_ids = input_ids[:, :input_seq_len]
+        prompts = [self.tokenizer.decode(query, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+                    for query in prompt_ids]
+        eos_probs = eos_logprobs.exp()
+
+        return {
+            'query/input_ids': prompt_ids,
+            'query/text': prompts,
+            'query/mask': attention_mask,
+            'response/input_ids': response_ids,
+            'response/text': response_text,
+            'response/mask': output_mask,
+            'response/log_prob': output_logprob,
+            'response/eos_prob': eos_probs,
+        }
+
+    def forward_pass(self,
+                     query_input_ids: torch.Tensor,
+                     query_mask: torch.Tensor,
+                     response_input_ids: torch.Tensor,
+                     response_mask: torch.Tensor,
+                     features: torch.Tensor,
+                     labels: Optional[torch.Tensor] = None,
+                     invalidate_eos: bool = True,
+                     device = None):
+
+        if device is None:
+            device = self.device
+
+        batch_size, query_seq_len = query_input_ids.shape
+        input_ids = torch.cat([query_input_ids, response_input_ids], dim=-1)
+        attention_mask = torch.cat([query_mask, response_mask], dim=-1)
+
+        # forward pass to get next token
+        outputs = self.model(
+            input_ids,
+            features,
+            attention_mask,
+            labels,
+            device=device
+        )
+        # get the first logit
+        query_logits = outputs.logits[:, :query_seq_len, :]
+        last_non_masked_idx = torch.sum(query_mask, dim=1) - 1
+        first_logits = query_logits[range(batch_size), last_non_masked_idx, :]
+        # get the second to last logit
+        response_logits = outputs.logits[:, query_seq_len:-1, :]
+        logits = torch.cat([first_logits[:, None], response_logits], dim=1)
+
+        negative_inf = HALF_NEGATIVE_INF if logits.dtype == torch.half else NEGATIVE_INF
+        if invalidate_eos:
+            logits[:, :, self.tokenizer.eos_token_id] = negative_inf  # no endoftext
+
+        log_prob = F.log_softmax(logits, dim=-1)
+        output_logprob = torch.gather(log_prob, 2, response_input_ids[:, :, None]).squeeze(2)
+        output_entropy = logits_to_entropy(logits)
+        eos_prob = F.softmax(logits, dim=-1)[:, :, self.tokenizer.eos_token_id]
+
+        pos_logit = torch.gather(logits, 2, response_input_ids[:, :, None]).squeeze(2)
+
+        return {
+            'response/log_prob': mask_pad(output_logprob, response_mask),
+            'response/eos_prob': mask_pad(eos_prob, response_mask),
+            'response/entropy': mask_pad(output_entropy, response_mask),
+            'response/pos_logit': mask_pad(pos_logit, response_mask),
+            'response/logits': logits,
+        }

requirements.txt

@@ -0,0 +1,19 @@
+matplotlib
+more-itertools
+pyyaml==5.4
+pillow
+numpy
+six
+tqdm
+ftfy
+regex
+huggingface-hub
+ipdb
+toml
+torch==1.11.0
+torchvision
+tensorboard
+transformers
+clip-anytorch==2.4.0
+gdown
+gradio

run.py

@@ -0,0 +1,173 @@
+import os
+import math
+import platform
+import logging
+from pathlib import Path
+
+import torch
+from transformers import AutoModelForCausalLM
+from PIL import Image
+import numpy as np
+from numpy import asarray
+import gradio as gr
+import clip
+
+from arguments import get_args
+from load import load_model_args, load_model
+from utils import get_first_sentence
+
+
+logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))
+log = logging.getLogger(__name__)
+
+
+def prepare(args):
+    num_gpus = torch.cuda.device_count()
+    log.info(f'Detect {num_gpus} GPUS')
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    args = load_model_args(args)
+
+    def load_style(args, checkpoint):
+        model = AutoModelForCausalLM.from_pretrained(args.init_model)
+        if checkpoint is not None and Path(checkpoint).is_file():
+            log.info("joint model: loading pretrained style generator")
+            state = torch.load(checkpoint)
+            if 'global_step' in state:
+                step = state['global_step']
+                log.info(f'trained for {step} steps')
+            weights = state['state_dict']
+            key = 'model.'
+            weights = {k[len(key):]: v for k, v in weights.items() if k.startswith(key)}
+            model.load_state_dict(weights)
+        else:
+            log.info("joint model: loading vanila gpt")
+        return model
+
+    log.info(f'loading models')
+    joint_model = load_style(args, checkpoint=getattr(args, 'demo_joint_model_weight', 'None'))
+    joint_model = joint_model.to(device)
+    model = load_model(args, device)
+    tokenizer = model.tokenizer
+    log.info(f'loaded models ')
+
+    class Inferer:
+        def __init__(self, args, model, joint_model, tokenizer, device):
+            self.args = args
+            self.model = model
+            self.joint_model = joint_model
+            self.tokenizer = tokenizer
+            self.device = device
+
+            self.clip_model, self.clip_preprocess = clip.load(args.clip_model_type, device=device, jit=False)
+
+        def infer_joint(self, batch, window_size=10, vanilla_length=20, sample=False, temperature=0.7, **kwargs):
+            with torch.no_grad():
+                rollouts = self.model.sample(input_ids=batch['input_ids'], attention_mask=batch['attention_mask'],
+                                        features=batch['features'], labels=None,
+                                        max_len=self.args.response_length, sample=sample,
+                                        no_repeat_ngram_size=self.args.infer_no_repeat_size,
+                                        invalidate_eos=False)
+                '''
+                query = rollouts['query/input_ids']
+                res = rollouts['response/input_ids']
+                gen1 = torch.cat([query, res], dim=1)
+                mask1 = torch.cat([rollouts['query/mask'], rollouts['response/mask']], dim=1)
+                '''
+                res = rollouts['response/text']
+                query = rollouts['query/text']
+                generations = [f'{q} {v.strip()}' for q, v in zip(query, res)]
+
+                cur_length = self.args.response_length
+                if vanilla_length > 0:
+                    for i in range(math.ceil(vanilla_length / window_size)):
+                        cur_length += window_size
+                        generations = self.tokenizer(generations, padding=True, return_tensors='pt').to(self.device)
+                        context = generations['input_ids'][:, :-window_size]
+                        inputs = generations['input_ids'][:, -window_size:]
+                        out = self.joint_model.generate(input_ids=inputs,
+                                                max_length=cur_length, sample=sample,
+                                                no_repeat_ngram_size=self.args.infer_no_repeat_size,
+                                                        pad_token_id=self.tokenizer.eos_token_id)
+                        out = torch.cat([context, out], dim=1)
+                        text = [self.tokenizer.decode(v, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+                                for v in out]
+                        # generations = [get_first_sentence(v) for v in generations]
+                        generations = text
+                query = rollouts['query/text']
+                del rollouts
+            torch.cuda.empty_cache()
+            return query, generations
+
+        def get_feature(self, image):
+            image = self.clip_preprocess(image).unsqueeze(0).to(self.device)
+            feature = self.clip_model.encode_image(image)
+            return feature
+
+        def __call__(self, image, prompt, length=20, window_size=20, **kwargs):
+            window_size = min(window_size, length)
+            vanilla_length = max(0, length - self.args.response_length)
+            if not prompt:
+                prompt = 'The'
+            feature = self.get_feature(image)
+            feature = feature.unsqueeze(0).to(self.device)
+            batch = self.tokenizer(prompt, padding=True, return_tensors='pt').to(self.device)
+            batch['features'] = feature
+            query, generations = self.infer_joint(batch, window_size=window_size,
+                                                  vanilla_length=vanilla_length, **kwargs)
+            # text = f'{query[0].strip()} {generations[0].strip()}'
+            text = generations[0].strip()
+            return text
+
+    inferer = Inferer(args, model, joint_model, tokenizer, device)
+    return inferer
+
+
+class Runner:
+    def __init__(self, inferer):
+        self.inferer = inferer
+
+    def __call__(self, inp, prompt, length, window_size, sample):
+        # inp = inp.reshape((224, 224, 3))
+        img = Image.fromarray(np.uint8(inp))
+        text = self.inferer(img, prompt, length, window_size, sample=sample)
+        return prompt, text
+        # return inp, prompt, text
+
+
+'''
+# test_run
+sample_img = asarray(Image.open('../data/coco/images/sample.jpg'))
+img, _, text = run(sample_img, 'There lies', 50, 20, sample=False)
+print('test_run:', text)
+'''
+
+def launch(examples=None):
+    args = get_args()
+    inferer = prepare(args)
+    runner = Runner(inferer)
+
+    iface = gr.Interface(
+        title="Demo for ESPER",
+        fn=runner.__call__,
+        inputs=[gr.components.Image(shape=(224, 224)),
+                gr.components.Textbox(label='prompt'),
+                gr.components.Slider(20, 120, step=1, label='length'),
+                gr.components.Slider(10, 100, step=1, label='window_size'),
+                gr.components.Checkbox(label='do sample')],
+        outputs=[gr.components.Textbox(label='prompt'),
+                gr.components.Textbox(label='generation')],
+        examples=examples
+    )
+    if args.port is not None:
+        print(f"running from {platform.node()}")
+        iface.launch(
+            server_name="0.0.0.0",
+            server_port=args.port
+        )
+    else:
+        iface.launch()
+
+
+if __name__ == "__main__":
+    print(f"running from {platform.node()}")
+    launch()

utils.py

@@ -0,0 +1,28 @@
+import torch
+
+
+NEGATIVE_INF = -100000.0
+HALF_NEGATIVE_INF = -60000.0  # half precision
+
+
+def get_first_sentence(txt, min_len=5):
+    eos = '<|endoftext|>'
+    eos_idx = txt.find(eos)
+    if eos_idx > 0:
+        txt = txt[eos_idx:]
+    txt = txt.replace('\n', ' ')
+    sents = txt.split('. ')
+    if len(sents[0]) >= min_len:
+        sent = f'{sents[0].strip()}.'
+    else:
+        sent = txt
+    return sent
+
+
+def logits_to_entropy(logits):
+    distribution = torch.distributions.Categorical(logits=logits)
+    return distribution.entropy()
+
+
+def mask_pad(value, mask):
+    return value * mask + NEGATIVE_INF * (1 - mask)