get_bert_encode函数代码

以下是一个使用Hugging Face Transformers库中的BERT模型进行编码的示例函数：

import torch
from transformers import BertTokenizer, BertModel

def get_bert_encode(text, model_name='bert-base-uncased', pooling_method='mean'):
    # Load pre-trained BERT model and tokenizer
    tokenizer = BertTokenizer.from_pretrained(model_name)
    model = BertModel.from_pretrained(model_name)

    # Tokenize input text
    tokenized_text = tokenizer.encode(text, add_special_tokens=True)

    # Convert tokenized_text to a tensor
    input_ids = torch.tensor([tokenized_text])

    # Get BERT model output
    with torch.no_grad():
        output = model(input_ids)

    # Extract the last hidden state of the BERT model
    last_hidden_state = output[0]

    # Apply pooling to get a fixed-length vector representation of the input text
    if pooling_method == 'mean':
        pooled_output = torch.mean(last_hidden_state, dim=1)
    elif pooling_method == 'max':
        pooled_output = torch.max(last_hidden_state, dim=1)[0]
    else:
        raise ValueError("Invalid pooling method. Must be either 'mean' or 'max'.")

    # Convert the tensor to a numpy array
    encoded_text = pooled_output.squeeze().numpy()

    return encoded_text

这个函数使用了BertTokenizer类来将输入文本转换为BERT模型可以接受的格式，并使用BertModel类进行编码。该函数还支持使用“平均”或“最大”池化方法将BERT的最后一个隐藏状态转换为固定长度的向量表示形式。