| |
| from models.proj_models import LLM2EmbedProjector, LLM2EmbedAddProjector |
| from typing import Optional |
| from logger import get_logger |
| from transformers.modeling_utils import load_state_dict |
| from transformers.trainer_utils import get_last_checkpoint |
| import os |
| from . import register_loader |
| logger = get_logger(__name__) |
|
|
|
|
| @register_loader("load_mlp_projector") |
| def load_mlp_projector( |
| fc_args: "FCArguments", |
| in_dim = 3074, |
| is_trainable: Optional[bool] = False, |
| resume_from_checkpoint=None, |
| output_dir=None |
| ): |
| if fc_args.stage == "classify": |
| mlp_model = LLM2EmbedProjector( |
| in_dim = in_dim, |
| out_dim = fc_args.embed_dim, |
| fp16=False |
| ) |
|
|
| if resume_from_checkpoint: |
| if os.path.exists(output_dir): |
| new_model_name_or_path_proj = get_last_checkpoint(output_dir) |
| logger.info("start to load pfc not from: {}, but from {}".format( |
| fc_args.model_name_or_path_proj, new_model_name_or_path_proj |
| )) |
| else: |
| new_model_name_or_path_proj = None |
| logger.info('not exists (output_dir) for pfc {}, load pfc from {}'.format(output_dir, fc_args.model_name_or_path_proj)) |
|
|
| if new_model_name_or_path_proj: |
| logger.info("(resume process) start to load pfc from: {}".format(new_model_name_or_path_proj)) |
| mlp_model.maybe_load_state_dict_from_path(new_model_name_or_path_proj) |
| else: |
| logger.info("(resume failed) still start to load pfc from: {}".format(fc_args.model_name_or_path_proj)) |
| mlp_model.maybe_load_state_dict_from_path(fc_args.model_name_or_path_proj) |
| else: |
| logger.info("start to load pfc from: {}".format(fc_args.model_name_or_path_proj)) |
| mlp_model.maybe_load_state_dict_from_path(fc_args.model_name_or_path_proj) |
|
|
|
|
| if is_trainable: |
| mlp_model.train() |
| else: |
| mlp_model.eval() |
| else: |
| mlp_model = None |
| return mlp_model |
|
|
|
|
| @register_loader("load_add_projector") |
| def load_add_projector( |
| fc_args: "FCArguments", |
| in_dim = 3074, |
| is_trainable: Optional[bool] = False, |
| resume_from_checkpoint=None, |
| output_dir=None |
| ): |
| if fc_args.stage == "classify": |
| mlp_model = LLM2EmbedAddProjector( |
| in_dim = in_dim, |
| out_dim = fc_args.embed_dim, |
| fp16=False |
| ) |
| logger.info("No weights for LLM2EmbedAddProjector") |
| if is_trainable: |
| mlp_model.train() |
| else: |
| mlp_model.eval() |
| else: |
| mlp_model = None |
| return mlp_model |
