1.安装必要的库
确保安装 transformers 和 torch 库:
pip install transformers torch
2.加载本地 BERT 模型和分词器
由于已将模型和分词器下载到本地,可以指定文件路径加载。确保路径与本地文件结构一致。
from transformers import BertTokenizer, BertModel
# 指定模型和分词器的路径
BERT_PATH = 'D:/bert/241109'
# 加载分词器和模型
tokenizer = BertTokenizer.from_pretrained(BERT_PATH, local_files_only=True)
model = BertModel.from_pretrained(BERT_PATH, local_files_only=True)
print("BERT 模型和分词器加载完成")
3.读取 .txt 文件内容
# 定义文件路径
file_path = 'D:/bert/AIsecKG-cybersecurity-dataset-main/datasource/textfiles/lab1.txt'
# 读取文件内容
with open(file_path, 'r', encoding='utf-8') as file:
text = file.read()
print("文本文件内容读取完成")
4.处理文本并获取 BERT 模型的输出
将读取的文本内容传入分词器并使用 BERT 模型进行处理。
# 使用分词器对文本进行编码,并自动截断超长的输入
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
# 获取模型输出
outputs = model(**inputs)
# 获取最后一层的隐藏状态(通常用于特征表示)
last_hidden_states = outputs.last_hidden_state
print("BERT 模型输出已获取")
5.输出或保存结果
可以根据任务需要进一步处理 last_hidden_states 或者 outputs。例如,可以提取池化输出用于文本分类或其他任务。
# 使用池化输出作为文本的整体表示
pooled_output = outputs.pooler_output
print("文本的整体表示:", pooled_output)