from typing import Callable, List, Optional, Tuple, Union from collections import namedtuple import json import glob import math import numpy as np import os import torch from torch import Tensor import torch.nn as nn import torch.nn.functional as F from einops import rearrange from functools import partial import pickle as pkl from PIL import Image, UnidentifiedImageError from transformers import AutoConfig, AutoModel, AutoModelForCausalLM from transformers import OPTForCausalLM, GPT2Tokenizer from transformers import CLIPVisionModel, CLIPVisionConfig from fromage import utils class FrozenArgs: freeze_lm: bool = True freeze_vm: bool = True opt_version: str = 'facebook/opt-6.7b' visual_encoder: str = 'openai/clip-vit-large-patch14' n_visual_tokens: int = 1 image_embed_dropout_prob: float = 0.0 task: str = 'captioning' shared_emb_dim: Optional[int] = 256 text_emb_layers: List[int] = [-1] retrieval_token_idx: int = 0 class FromageModel(nn.Module): def __init__(self, tokenizer, args: FrozenArgs = FrozenArgs()): super().__init__() self.tokenizer = tokenizer self.feature_extractor = utils.get_feature_extractor_for_model(args.visual_encoder, train=False) self.image_token = self.tokenizer.cls_token_id assert args.text_emb_layers != set(args.text_emb_layers), 'text_emb_layers not unique' self.args = args opt_version = args.opt_version visual_encoder = args.visual_encoder n_visual_tokens = args.n_visual_tokens print(f"Using {opt_version} for the language model.") print(f"Using {visual_encoder} for the visual model with {n_visual_tokens} visual tokens.") self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) if 'facebook/opt' in opt_version: self.lm = OPTForCausalLM.from_pretrained(opt_version) else: raise NotImplementedError self.opt_version = opt_version if self.args.freeze_lm: self.lm.eval() print("Freezing the LM.") for param in self.lm.parameters(): param.requires_grad = False else: self.lm.train() self.retrieval_token_idx = args.retrieval_token_idx print(f'Initializing embedding for the retrieval token [RET] (id = {self.retrieval_token_idx}).') self.lm.resize_token_embeddings(len(tokenizer)) self.input_embeddings = self.lm.get_input_embeddings() print("Restoring pretrained weights for the visual model.") if 'clip' in visual_encoder: self.visual_model = CLIPVisionModel.from_pretrained(visual_encoder) else: self.visual_model = AutoModel.from_pretrained(visual_encoder) if 'clip' in visual_encoder: hidden_size = self.visual_model.config.hidden_size else: raise NotImplementedError if self.args.freeze_vm: print("Freezing the VM.") self.visual_model.eval() for param in self.visual_model.parameters(): param.requires_grad = False else: self.visual_model.train() self.visual_model_name = visual_encoder embedding_dim = self.input_embeddings.embedding_dim * self.args.n_visual_tokens self.text_hidden_fcs = nn.ModuleList([]) if self.args.shared_emb_dim is None: if len(self.args.text_emb_layers) == 1: if (self.args.text_emb_layers[0] in [-1, self.lm.config.num_hidden_layers]) and ('bert' not in opt_version): out_dim = self.lm.config.word_embed_proj_dim else: out_dim = self.lm.config.hidden_size else: if (-1 in self.args.text_emb_layers) or (self.lm.config.num_hidden_layers in self.args.text_emb_layers) \ and (self.lm.config.word_embed_proj_dim != self.lm.config.hidden_size): raise ValueError('No projection dim specified but model uses last output layer and an intermediate one (which have different dims).') else: out_dim = self.lm.config.hidden_size else: out_dim = self.args.shared_emb_dim for layer_idx in self.args.text_emb_layers: if (layer_idx == -1 or layer_idx == self.lm.config.num_hidden_layers) and ('bert' not in opt_version): in_dim = self.lm.config.word_embed_proj_dim text_fc = [nn.Linear(in_dim, out_dim), nn.Dropout(self.args.text_embed_dropout_prob)] self.text_hidden_fcs.append(nn.Sequential(*text_fc)) elif layer_idx < self.lm.config.num_hidden_layers: text_fc = [nn.Linear(self.lm.config.hidden_size, out_dim), nn.Dropout(self.args.text_embed_dropout_prob)] self.text_hidden_fcs.append(nn.Sequential(*text_fc)) else: raise ValueError(f'Embedding of layer {layer_idx} was requested but model only has {self.lm.config.num_hidden_layers} layers.') self.visual_embeddings = nn.Linear(hidden_size, embedding_dim) self.visual_fc = nn.Linear(hidden_size, out_dim) self.image_dropout = nn.Dropout(self.args.image_embed_dropout_prob) def get_visual_embs(self, pixel_values: torch.FloatTensor, mode: str = 'captioning'): if mode not in ['captioning', 'retrieval']: raise ValueError(f'mode should be one of ["caption", "retrieval"], got {mode} instead.') # Extract visual embeddings from the vision encoder. if 'clip' in self.visual_model_name: outputs = self.visual_model(pixel_values) encoder_outputs = outputs.pooler_output else: raise NotImplementedError # Use the correct fc based on function argument. if mode == 'captioning': visual_embs = self.visual_embeddings(encoder_outputs) # (2, D * n_visual_tokens) visual_embs = torch.reshape(visual_embs, (visual_embs.shape[0], self.args.n_visual_tokens, -1)) elif mode == 'retrieval': visual_embs = self.visual_fc(encoder_outputs) # (2, D * n_visual_tokens) visual_embs = torch.reshape(visual_embs, (visual_embs.shape[0], 1, -1)) else: raise NotImplementedError visual_embs = self.image_dropout(visual_embs) return visual_embs def train(self, mode=True): super(FromageModel, self).train(mode=mode) # Overwrite train() to ensure Frozen models remain frozen. if self.args.freeze_lm: self.lm.eval() if self.args.freeze_vm: self.visual_model.eval() def forward( self, pixel_values: torch.FloatTensor, labels: torch.LongTensor, caption_len: torch.LongTensor, mode: str = 'captioning', concat_captions: bool = False, input_prefix: Optional[str] = None, inference: bool = False, ): visual_embs = self.get_visual_embs(pixel_values, mode) batch_size, vis_seq_len, _ = visual_embs.shape # vis_seq_len = n_visual_tokens if labels is not None: assert labels.shape[0] == batch_size, (visual_embs.shape, labels.shape) input_embs = self.input_embeddings(labels) # (N, T, D) last_embedding_idx = caption_len - 1 # -1 to retrieve the token before the eos token if input_prefix is not None: prompt_ids = self.tokenizer(input_prefix, add_special_tokens=False, return_tensors="pt").input_ids prompt_ids = prompt_ids.to(visual_embs.device) prompt_embs = self.input_embeddings(prompt_ids) prompt_embs = prompt_embs.repeat(batch_size, 1, 1) assert prompt_embs.shape[0] == batch_size, prompt_embs.shape assert prompt_embs.shape[2] == input_embs.shape[2], prompt_embs.shape assert len(prompt_embs.shape) == 3, prompt_embs.shape if mode == 'captioning': # Concat to text embeddings. condition_seq_len = 0 if input_prefix is None: # Just add visual embeddings. input_embs = torch.cat([visual_embs, input_embs], axis=1) last_embedding_idx += vis_seq_len condition_seq_len += vis_seq_len full_labels = torch.zeros(visual_embs.shape[:2], dtype=torch.int64).to(visual_embs.device) - 100 else: # Add visual and prompt embeddings. prefix_embs = torch.cat([visual_embs, prompt_embs], axis=1) input_embs = torch.cat([prefix_embs, input_embs], axis=1) last_embedding_idx += prefix_embs.shape[1] condition_seq_len += prefix_embs.shape[1] full_labels = torch.zeros(prefix_embs.shape[:2], dtype=torch.int64).to(visual_embs.device) - 100 # Mask out embedding tokens in the labels. full_labels = torch.cat([full_labels, labels], axis=1) pad_idx = [] for label in full_labels: for k, token in enumerate(label): # Mask out retrieval token if it exists. if token in [self.tokenizer.pad_token_id, self.retrieval_token_idx]: label[k:] = -100 pad_idx.append(k) break if k == len(label) - 1: # No padding found. pad_idx.append(k + 1) assert len(pad_idx) == batch_size, (len(pad_idx), batch_size) bs, seq_len, embs_dim = input_embs.shape if concat_captions: assert len(input_embs.shape) == 3, input_embs assert len(full_labels.shape) == 2, full_labels assert batch_size % 2 == 0 all_concat_input_embs = [] all_concat_labels = [] # Rearrange embeddings and labels (and their padding) to concatenate captions. for i in range(batch_size // 2): first_idx = i * 2 second_idx = first_idx + 1 first_emb = input_embs[first_idx, :pad_idx[first_idx], :] first_labels = full_labels[first_idx, :pad_idx[first_idx]] first_padding = input_embs[first_idx, pad_idx[first_idx]:, :] first_labels_padding = full_labels[first_idx, pad_idx[first_idx]:] second_emb = input_embs[second_idx, :pad_idx[second_idx], :] second_labels = full_labels[second_idx, :pad_idx[second_idx]] second_padding = input_embs[second_idx, pad_idx[second_idx]:, :] second_labels_padding = full_labels[second_idx, pad_idx[second_idx]:] assert torch.all(first_labels_padding == -100), first_labels_padding assert torch.all(second_labels_padding == -100), second_labels_padding concat_input_embs = torch.cat([first_emb, second_emb, first_padding, second_padding], axis=0) # (T*2, 768) concat_labels = torch.cat([first_labels, second_labels, first_labels_padding, second_labels_padding], axis=0) # (T*2, 768) all_concat_input_embs.append(concat_input_embs) all_concat_labels.append(concat_labels) # Pad to max length. input_embs = torch.stack(all_concat_input_embs, axis=0) # (N/2, T*2, 768) full_labels = torch.stack(all_concat_labels, axis=0) # (N/2, T*2, 768) assert input_embs.shape == (bs // 2, seq_len * 2, embs_dim), input_embs.shape assert full_labels.shape == (bs // 2, seq_len * 2), full_labels.shape output = self.lm(inputs_embeds=input_embs, labels=full_labels, output_hidden_states=True) elif mode == 'retrieval': full_labels = torch.clone(labels) if input_prefix is not None: print(f'Adding prefix "{input_prefix}" to retrieval.') # Add prompt embeddings. prefix_embs = prompt_embs input_embs = torch.cat([prefix_embs, input_embs], axis=1) last_embedding_idx += prefix_embs.shape[1] full_labels = torch.cat([ torch.zeros(prefix_embs.shape[:2], dtype=torch.int64).to(labels.device) - 100, full_labels ], axis=1) pad_idx = [] for label in full_labels: for k, token in enumerate(label): if token == self.tokenizer.pad_token_id: label[k:] = -100 pad_idx.append(k) break if k == len(label) - 1: # No padding found. pad_idx.append(k + 1) assert len(pad_idx) == batch_size, (len(pad_idx), batch_size) output = self.lm(inputs_embeds=input_embs, labels=full_labels, output_hidden_states=True) else: raise NotImplementedError last_embedding = None last_output_logit = None hidden_states = [] if mode == 'retrieval': if self.args.shared_emb_dim is not None: for idx, fc_layer in zip(self.args.text_emb_layers, self.text_hidden_fcs): hidden_states.append(fc_layer(output.hidden_states[idx])) # (N, seq_len, 2048) else: for idx in self.args.text_emb_layers: hidden_states.append(output.hidden_states[idx]) # Add hidden states together. last_hidden_state = torch.stack(hidden_states, dim=-1).sum(dim=-1) if not concat_captions: last_embedding = torch.stack([last_hidden_state[i, last_embedding_idx[i], :] for i in range(batch_size)], axis=0) # (N, D) last_output_logit = torch.stack([output.logits[i, last_embedding_idx[i] - 1, :] for i in range(batch_size)], axis=0) # (N, D) else: # Concatenate two captioning examples together. all_last_embedding = [] all_last_output_logit = [] for i in range(batch_size // 2): first_last_embedding_idx, second_last_embedding_idx = all_last_embedding_idx[i] first_last_embedding = last_hidden_state[i, first_last_embedding_idx, :] # (N, D) first_last_output_logit = output.logits[i, first_last_embedding_idx - 1, :] # (N, D) second_last_embedding = last_hidden_state[i, second_last_embedding_idx, :] # (N, D) second_last_output_logit = output.logits[i, second_last_embedding_idx - 1, :] # (N, D) all_last_embedding.append(first_last_embedding) all_last_embedding.append(second_last_embedding) all_last_output_logit.append(first_last_output_logit) all_last_output_logit.append(second_last_output_logit) last_embedding = torch.stack(all_last_embedding) last_output_logit = torch.stack(all_last_output_logit) # Compute retrieval loss. assert visual_embs.shape[1] == 1, visual_embs.shape visual_embs = visual_embs[:, 0, :] visual_embs = visual_embs / visual_embs.norm(dim=1, keepdim=True) last_embedding = last_embedding / last_embedding.norm(dim=1, keepdim=True) # cosine similarity as logits logit_scale = self.logit_scale.exp() visual_embs = logit_scale * visual_embs elif mode == 'captioning': pass else: raise NotImplementedError return output, full_labels, last_embedding, last_output_logit, visual_embs def generate(self, embeddings = torch.FloatTensor, max_len: int = 32, temperature: float = 0.0, top_p: float = 1.0, min_word_tokens: int = 0, ret_scale_factor: float = 1.0, filter_value: float = -float('Inf')): """Runs greedy decoding and returns generated captions. Args: embeddings: Input condition that the model uses for autoregressive generation. max_len: Maximum number of tokens to generate. temperature: Used to modulate logit distribution. top_p: If set to < 1, the smallest set of tokens with highest probabilities that add up to top_p or higher are kept for generation. min_word_tokens: Minimum number of words to generate before allowing a [RET] output. ret_scale_factor: Proportion to scale [RET] token logits by. A higher value may increase the probability of the model generating [RET] outputs. filter_value: Value to assign to tokens that should never be generated. Outputs: out: (N, T) int32 sequence of output tokens. output_embeddings: (N, T, 256) sequence of text output embeddings. """ self.lm.eval() with torch.no_grad(): # no tracking history batch_size, s, _ = embeddings.shape # init output with image tokens out = None past_key_values = None output_embeddings = [] output_logits = [] for i in range(max_len): if 'opt' in self.opt_version: output = self.lm(inputs_embeds=embeddings, use_cache=False, output_hidden_states=True) else: if i == 0: output = self.lm(inputs_embeds=embeddings, use_cache=True, past_key_values=None, output_hidden_states=True) else: output = self.lm(input_ids=out[:, -1:], use_cache=True, past_key_values=past_key_values, output_hidden_states=True) # Collect and sum the hidden states. hidden_states = [] if self.args.shared_emb_dim is not None: for idx, fc_layer in zip(self.args.text_emb_layers, self.text_hidden_fcs): hidden_states.append(fc_layer(output.hidden_states[idx])) # (N, seq_len, 2048) else: for idx in self.args.text_emb_layers: hidden_states.append(output.hidden_states[idx]) # Add hidden states together. last_hidden_state = torch.stack(hidden_states, dim=-1).sum(dim=-1) # (N, T, 256) last_embedding = last_hidden_state / last_hidden_state.norm(dim=-1, keepdim=True) output_embeddings.append(last_embedding) logits = output.logits[:, -1, :] # (N, vocab_size) if top_p == 1.0: logits = logits.cpu() output_logits.append(logits) if self.retrieval_token_idx != -1 and self.retrieval_token_idx is not None: if i < min_word_tokens: # Eliminate probability of generating [RET] if this is earlier than min_word_tokens. logits[:, self.retrieval_token_idx] = filter_value else: # Multiply by scaling factor. logits[:, self.retrieval_token_idx] = logits[:, self.retrieval_token_idx] * ret_scale_factor past_key_values = output.past_key_values if temperature == 0.0: if top_p != 1.0: raise ValueError('top_p cannot be set if temperature is 0 (greedy decoding).') next_token = torch.argmax(logits, keepdim=True, dim=-1) # (N, 1) else: logits = logits / temperature # Apply top-p filtering. if top_p < 1.0: assert top_p > 0, f'top_p should be above 0, got {top_p} instead.' sorted_logits, sorted_indices = torch.sort(logits, descending=True) # (N, D) and (N, D) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) # (N, D) # Remove tokens with cumulative probability above the threshold sorted_indices_to_remove = cumulative_probs > top_p # Shift the indices to the right to keep also the first token above the threshold sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] = 0 for j in range(sorted_indices.shape[0]): indices_to_remove = sorted_indices[j, sorted_indices_to_remove[j, :]] logits[j, indices_to_remove] = filter_value token_weights = logits.exp() # (N, vocab_size) next_token = torch.multinomial(token_weights, 1) # (N, 1) next_token = next_token.long().to(embeddings.device) if out is not None: out = torch.cat([out, next_token], dim=-1) else: out = next_token if 'opt' in self.opt_version: next_embedding = self.input_embeddings(next_token) embeddings = torch.cat([embeddings, next_embedding], dim=1) elif (self.tokenizer.eos_token_id and (next_token == self.tokenizer.eos_token_id).all()): # End of generation. break return out, output_embeddings, output_logits class Fromage(nn.Module): def __init__(self, tokenizer, model_args: Optional[FrozenArgs] = None, path_array: Optional[List[str]] = None, emb_matrix: Optional[torch.tensor] = None): super().__init__() self.model = FromageModel(tokenizer, model_args) self.path_array = path_array self.emb_matrix = emb_matrix def __call__(self, images: Tensor, tgt_tokens: Optional[Tensor] = None, caption_len: Optional[Tensor] = None, generate: bool = False, num_words: int = 32, temperature: float = 1.0, top_p: float = 1.0, ret_scale_factor: float = 1.0, min_word_tokens: int = 0, mode: str = 'captioning', concat_captions: bool = False, input_prefix: Optional[str] = None, inference: bool = False) -> Tensor: if generate: return self.model.generate(images, num_words, temperature=temperature, top_p=top_p, min_word_tokens=min_word_tokens, ret_scale_factor=ret_scale_factor) else: output = self.model( pixel_values = images, labels = tgt_tokens, caption_len = caption_len, mode = mode, concat_captions = concat_captions, input_prefix = input_prefix, inference = inference) return output def generate_for_images_and_texts( self, prompts: List, num_words: int = 0, ret_scale_factor: float = 1.0, top_p: float = 1.0, temperature: float = 0.0, max_num_rets: int = 1): """ Encode prompts into embeddings. Args: prompts: List of interleaved PIL.Image.Image and strings representing input to the model. num_words: Maximum number of words to generate for. If num_words = 0, the model will run its forward pass and return the outputs. ret_scale_factor: Proportion to scale [RET] token logits by. A higher value may increase the probability of the model generating [RET] outputs. top_p: If set to < 1, the smallest set of tokens with highest probabilities that add up to top_p or higher are kept for generation. temperature: Used to modulate logit distribution. max_num_rets: Maximum number of images to return in one generation pass. Returns: return_outputs: List consisting of either str or List[PIL.Image.Image] objects, representing image-text interleaved model outputs. """ input_embs = [] input_ids = [] add_bos = True for i, p in enumerate(prompts): if type(p) == Image.Image: # Encode as image. pixel_values = utils.get_pixel_values_for_model(self.model.feature_extractor, p) pixel_values = pixel_values.to(device=self.model.logit_scale.device, dtype=self.model.logit_scale.dtype) pixel_values = pixel_values[None, ...] visual_embs = self.model.get_visual_embs(pixel_values, mode='captioning') # (1, n_visual_tokens, D) input_embs.append(visual_embs) elif type(p) == str: text_ids = self.model.tokenizer(p, add_special_tokens=True, return_tensors="pt").input_ids.to(self.model.logit_scale.device) if not add_bos: # Remove tag. text_ids = text_ids[:, 1:] else: # Only add once. add_bos = False text_embs = self.model.input_embeddings(text_ids) # (1, T, D) input_embs.append(text_embs) input_ids.append(text_ids) else: raise ValueError(f'Input prompts should be either PIL.Image.Image or str types, got {type(p)} instead.') input_embs = torch.cat(input_embs, dim=1) input_ids = torch.cat(input_ids, dim=1) if num_words == 0: generated_ids = input_ids outputs = self.model.lm(inputs_embeds=input_embs, use_cache=False, output_hidden_states=True) # Map outputs to embeddings, so we can retrieve embeddings from the [RET] tokens. out = [] for x, fc in zip(self.model.args.text_emb_layers, self.model.text_hidden_fcs): out.append(fc(outputs.hidden_states[x])) embeddings = torch.stack(out, dim=-1).sum(dim=-1) embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True) # (N, T, 256) elif num_words > 0: generated_ids, generated_embeddings, _ = self.model.generate(input_embs, num_words, temperature=temperature, top_p=top_p, ret_scale_factor=ret_scale_factor) embeddings = generated_embeddings[-1][:, input_embs.shape[1]:] # Truncate to newline. newline_token_id = self.model.tokenizer('\n', add_special_tokens=False).input_ids[0] trunc_idx = 0 for j in range(generated_ids.shape[1]): if generated_ids[0, j] == newline_token_id: trunc_idx = j break if trunc_idx > 0: generated_ids = generated_ids[:, :trunc_idx] embeddings = embeddings[:, :trunc_idx] else: raise ValueError # Save outputs as an interleaved list. return_outputs = [] # Find up to max_num_rets [RET] tokens, and their corresponding scores. all_ret_idx = [i for i, x in enumerate(generated_ids[0, :] == self.model.retrieval_token_idx) if x][:max_num_rets] seen_image_idx = [] # Avoid showing the same image multiple times. last_ret_idx = 0 if len(all_ret_idx) == 0: # No [RET] tokens. caption = self.model.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] return_outputs.append(utils.truncate_caption(caption)) else: for ret_idx in all_ret_idx: ret_emb = embeddings[:, ret_idx, :] scores = self.emb_matrix @ ret_emb.T # Downweight seen images. for seen_idx in seen_image_idx: scores[seen_idx, :] -= 1000 # Get the top 3 images for each image. _, top_image_idx = scores.squeeze().topk(3) image_outputs = [] for img_idx in top_image_idx: # Find the first image that does not error out. try: seen_image_idx.append(img_idx) img = utils.get_image_from_url(self.path_array[img_idx]) image_outputs.append(img) if len(image_outputs) == max_num_rets: break except UnidentifiedImageError: pass caption = self.model.tokenizer.batch_decode(generated_ids[:, last_ret_idx:ret_idx], skip_special_tokens=True)[0] last_ret_idx = ret_idx + 1 return_outputs.append(utils.truncate_caption(caption) + ' [RET]') return_outputs.append(image_outputs) return return_outputs def load_fromage(model_dir: str, ckpt_path: str) -> Fromage: model_args_path = os.path.join(model_dir, 'model_args.json') model_ckpt_path = os.path.join(ckpt_path) embs_paths = [s for s in glob.glob(os.path.join(model_dir, 'cc3m_embeddings*.pkl'))] if not os.path.exists(model_args_path): raise ValueError(f'model_args.json does not exist in {model_dir}.') if not os.path.exists(model_ckpt_path): raise ValueError(f'pretrained_ckpt.pth.tar does not exist in {model_dir}.') if len(embs_paths) == 0: raise ValueError(f'cc3m_embeddings_*.pkl files do not exist in {model_dir}.') # Load embeddings. # Construct embedding matrix for nearest neighbor lookup. path_array = [] emb_matrix = [] # These were precomputed for all CC3M images with `model.get_visual_embs(image, mode='retrieval')`. for p in embs_paths: with open(p, 'rb') as wf: train_embs_data = pkl.load(wf) path_array.extend(train_embs_data['paths']) emb_matrix.append(train_embs_data['embeddings']) emb_matrix = np.concatenate(emb_matrix, axis=0) # Number of paths should be equal to number of embeddings. assert len(path_array) == emb_matrix.shape[0], (len(path_array), emb_matrix.shape[0]) with open(model_args_path, 'r') as f: model_kwargs = json.load(f) # Initialize tokenizer. tokenizer = GPT2Tokenizer.from_pretrained(model_kwargs['opt_version']) tokenizer.pad_token = tokenizer.eos_token # Add special tokens to the model to enable [RET]. tokenizer.add_special_tokens({"cls_token": "<|image|>"}) tokenizer.add_tokens('[RET]') ret_token_idx = tokenizer('[RET]', add_special_tokens=False).input_ids assert len(ret_token_idx) == 1, ret_token_idx model_kwargs['retrieval_token_idx'] = ret_token_idx[0] args = namedtuple('args', model_kwargs)(**model_kwargs) # Initialize model for inference. model = Fromage(tokenizer, args, path_array=path_array, emb_matrix=emb_matrix) model = model.eval() model = model.bfloat16() model = model.cuda() # Load pretrained linear mappings and [RET] embeddings. checkpoint = torch.load(model_ckpt_path) model.load_state_dict(checkpoint['state_dict'], strict=False) with torch.no_grad(): model.model.input_embeddings.weight[model.model.retrieval_token_idx, :].copy_(checkpoint['state_dict']['ret_input_embeddings.weight'].cpu().detach()) logit_scale = model.model.logit_scale.exp() emb_matrix = torch.tensor(emb_matrix, dtype=logit_scale.dtype).to(logit_scale.device) emb_matrix = emb_matrix / emb_matrix.norm(dim=1, keepdim=True) emb_matrix = logit_scale * emb_matrix model.emb_matrix = emb_matrix return model