autotrain로 fine tuning 후 embedding vector 구하는 방법이 있을까요?

안녕하세요.

mistral 과 llama2를 사용해서 embedding vector를 통해 챗봇을 구현하고 있습니다.

기존 모델로 embedding vector를 구하면 답변 검색이 잘 되지 않아서 파인튜닝을 한 후에 embedding vector를 구하려고 합니다.

학습과정에서 알려주신 대로 mistral 과 llama2의 fine tuning을 완료 했습니다.
fine tuning한 모델에서 generate는 잘 동작하나, embedding vector가 생성이 되질 않아 질문 드립니다.

기존 모델을 통한 임베딩 방식은 다음과 같습니다.

seq_ids = tokenizer(text, return_tensors='pt')["input_ids"]

embedding = model(seq_ids)["last_hidden_state"].mean(axis=[0,1]).detach().numpy()

기존 원본 모델에서는 'last_hidden_state' 값을 통해서 계산하지만

fine tuning한 모델에서는 'logits' 값만 존재 합니다. 

- 원본 모델 리턴값 :             odict_keys(['last_hidden_state'])

- fine tuning 모델 리턴값 : odict_keys(['logits'])

그래서 파인튜닝한 모델을 보면 Peft 와 Lora로 한번 레이어로 감싼 형태로 리턴되서 그럴꺼라 추측하는데요.

기존 모델

MistralForCausalLM(
  (model): MistralModel(
    (embed_tokens): Embedding(46081, 4096, padding_idx=2)
    (layers): ModuleList(
      (0-31): 32 x MistralDecoderLayer(
        (self_attn): MistralAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): MistralRotaryEmbedding()
        )
        (mlp): MistralMLP(
          (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
          (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
          (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): MistralRMSNorm()
        (post_attention_layernorm): MistralRMSNorm()
      )
    )
    (norm): MistralRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=46081, bias=False)
)

파인튜닝한 모델

PeftModelForCausalLM(
  (base_model): LoraModel(
    (model): MistralForCausalLM(
      (model): MistralModel(
        (embed_tokens): Embedding(46081, 4096, padding_idx=2)
        (layers): ModuleList(
          (0-31): 32 x MistralDecoderLayer(
            (self_attn): MistralAttention(
              (q_proj): Linear4bit(
                in_features=4096, out_features=4096, bias=False
                (lora_dropout): ModuleDict(
                  (default): Dropout(p=0.05, inplace=False)
                )
                (lora_A): ModuleDict(
                  (default): Linear(in_features=4096, out_features=16, bias=False)
                )
                (lora_B): ModuleDict(
                  (default): Linear(in_features=16, out_features=4096, bias=False)
                )
                (lora_embedding_A): ParameterDict()
                (lora_embedding_B): ParameterDict()
              )
              (k_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)
              (v_proj): Linear4bit(
                in_features=4096, out_features=1024, bias=False
                (lora_dropout): ModuleDict(
                  (default): Dropout(p=0.05, inplace=False)
                )
                (lora_A): ModuleDict(
                  (default): Linear(in_features=4096, out_features=16, bias=False)
                )
                (lora_B): ModuleDict(
                  (default): Linear(in_features=16, out_features=1024, bias=False)
                )
                (lora_embedding_A): ParameterDict()
                (lora_embedding_B): ParameterDict()
              )
              (o_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)
              (rotary_emb): MistralRotaryEmbedding()
            )
            (mlp): MistralMLP(
              (gate_proj): Linear4bit(in_features=4096, out_features=14336, bias=False)
              (up_proj): Linear4bit(in_features=4096, out_features=14336, bias=False)
              (down_proj): Linear4bit(in_features=14336, out_features=4096, bias=False)
              (act_fn): SiLUActivation()
            )
            (input_layernorm): MistralRMSNorm()
            (post_attention_layernorm): MistralRMSNorm()
          )
        )
        (norm): MistralRMSNorm()
      )
      (lm_head): Linear(in_features=4096, out_features=46081, bias=False)
    )
  )
)

fine tuning을 한 모델에서는 embedding vector 를 구하는게 불가능한건가요?

+추가 model.merge_and_unload() 를 써봤지만 소용없었습니다.

model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map={"":0}, output_hidden_states=True)