import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer, BertModel, BertTokenizer
from model.llmkt import *

def get_model_output_with_special_token(text, special_token="[SPECIAL]"):
    tokenizer, model = load_bert_model()
    tokenizer.add_special_tokens({"additional_special_tokens": [special_token]})
    model.resize_token_embeddings(len(tokenizer))
    
    text_with_special_token = text + " " + special_token
    inputs = tokenizer(text_with_special_token, return_tensors="pt")
    outputs = model(**inputs)
    
    # 打印输入和输出的维度
    print(f"Input shape: {inputs['input_ids'].shape}")
    print(f"Output shape: {outputs.last_hidden_state.shape}")
    
    # 获取特殊token对应位置的输出
    special_token_id = tokenizer.convert_tokens_to_ids(special_token)
    special_token_index = (inputs["input_ids"] == special_token_id).nonzero(as_tuple=True)
    special_token_output = outputs.last_hidden_state[special_token_index]
    
    print(f"Special token output: {special_token_output}")
    return outputs.last_hidden_state


