muxingyin/VisualGLM-6B

1import torch2from sat.model.official import ChatGLMModel3from sat.model.base_model import BaseMixin4from copy import deepcopy5import json6from .blip2 import BLIP27 8from sat.resources.urls import MODEL_URLS9MODEL_URLS['visualglm-6b'] = 'https://cloud.tsinghua.edu.cn/f/348b98dffcc940b6a09d/?dl=1'10 11class ImageMixin(BaseMixin):12    def __init__(self, args):13        super().__init__()14        self.args = deepcopy(args)15        self.model = BLIP2(args.eva_args, args.qformer_args)16 17    def word_embedding_forward(self, input_ids, output_cross_layer, **kw_args):18        if kw_args["pre_image"] > input_ids.shape[1] or kw_args.get("image", None) is None:19            return self.transformer.word_embeddings(input_ids)20        image_emb = self.model(**kw_args)21        # the image is inserted after 问：<img>, override 32 pads22        pre_id, pads, post_id = torch.tensor_split(input_ids, [kw_args["pre_image"], kw_args["pre_image"]+self.args.image_length], dim=1)23        pre_txt_emb = self.transformer.word_embeddings(pre_id)24        post_txt_emb = self.transformer.word_embeddings(post_id)25        return torch.cat([pre_txt_emb, image_emb, post_txt_emb], dim=1)26 27class VisualGLMModel(ChatGLMModel):28    def __init__(self, args, transformer=None, **kwargs):29        super().__init__(args, transformer=transformer, **kwargs)30        self.image_length = args.image_length31        self.add_mixin("eva", ImageMixin(args))32 33    @classmethod34    def add_model_specific_args(cls, parser):35        group = parser.add_argument_group('VisualGLM', 'VisualGLM Configurations')36        group.add_argument('--image_length', type=int, default=32)37        group.add_argument('--eva_args', type=json.loads, default={})38        group.add_argument('--qformer_args', type=json.loads, default={})39        return super().add_model_specific_args(parser)40    41