[大语言模型-工程实践] 手把手教你-基于BERT模型提取商品标题关键词及优化改进

pip instal transformers torch

from transformers import BertTokenizer, BertModel
import torch
# 确保你的设备可以运行BERT模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 加载预训练的BERT模型和分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')
model.to(device)
def extract_keywords(title, num_keywords=5):
    # 对标题进行分词
    tokens = tokenizer.tokenize(title)
    # 将分词结果转换为BERT模型可以处理的格式
    tokens = ['[CLS]'] + tokens + ['[SEP]']
    token_ids = tokenizer.convert_tokens_to_ids(tokens)
    token_type_ids = [0] * len(token_ids)  # 单句子输入，所有token_type_ids为0
    attention_mask = [1] * len(token_ids)  # 标记每个token是实际的输入
    # 将数据发送到设备
    token_ids = torch.tensor([token_ids]).to(device)
    token_type_ids = torch.tensor([token_type_ids]).to(device)
    attention_mask = torch.tensor([attention_mask]).to(device)
    # 模型预测
    with torch.no_grad():
        outputs = model(input_ids=token_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
    # 获取每个token的词向量
    embeddings = outputs[0]  # (batch_size, sequence_length, hidden_states)
    embeddings = embeddings.squeeze(0)  # 移除batch_size维度
    # 忽略特殊标记CLS和SEP
    embeddings = embeddings[1:-1]
    tokens = tokens[1:-1]
    # 基于词向量的平均值选择关键词
    keyword_scores = torch.mean(embeddings, dim=1)  # (sequence_length - 2, hidden_states)
    sorted_scores, sorted_indices = torch.sort(keyword_scores, descending=True, dim=0)
    # 获取最高分的词
    keywords_indices = sorted_indices[:num_keywords].cpu().numpy()
    keywords = [tokens[index] for index in keywords_indices]
    return keywords
# 示例标题
title = "2024 New Arrival! Women's Fashion Summer Dress"
keywords = extract_keywords(title, num_keywords=5)
print("Extracted Keywords:", keywords)

from transformers import BertTokenizer, BertForSequenceClassification, AdamW
from torch.utils.data import Dataset, DataLoader
class KeywordExtractionDataset(Dataset):
    def __init__(self, titles, keywords, tokenizer, max_len=128):
        self.titles = titles
        self.keywords = keywords
        self.tokenizer = tokenizer
        self.max_len = max_len
    def __len__(self):
        return len(self.titles)
    def __getitem__(self, idx):
        title = self.titles[idx]
        keyword = self.keywords[idx]
        inputs = self.tokenizer.encode_plus(
            title,
            None,
            add_special_tokens=True,
            max_length=self.max_len,
            padding='max_length',
            return_token_type_ids=True,
            truncation=True
        )
        inputs['labels'] = 1 if keyword else 0  # Simplified example
        return inputs
# 假设 titles 和 keywords 是已经准备好的数据列表
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
dataset = KeywordExtractionDataset(titles, keywords, tokenizer)
dataloader = DataLoader(dataset, batch_size=16, shuffle=True)
optimizer = AdamW(model.parameters(), lr=2e-5)
for epoch in range(num_epochs):
    model.train()
    for batch in dataloader:
        inputs = {k: v.to(device) for k, v in batch.items() if k != 'labels'}
        inputs['labels'] = batch['labels'].to(device)
        outputs = model(**inputs)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

import torch
import torch.nn as nn
class LoRALayer(nn.Module):
    def __init__(self, input_dim, output_dim, rank=4):
        super(LoRALayer, self).__init__()
        self.rank = rank
        self.A = nn.Parameter(torch.randn(output_dim, rank))
        self.B = nn.Parameter(torch.randn(rank, input_dim))
    def forward(self, x):
        return torch.matmul(self.A, torch.matmul(x, self.B).transpose(-1, -2))

from transformers import BertModel, BertConfig
class BertModelWithLoRA(BertModel):
    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.lora_layers = nn.ModuleDict()
        for i, layer in enumerate(self.encoder.layer):
            self.lora_layers[f"layer_{i}"] = LoRALayer(config.hidden_size, config.hidden_size, rank=4)
    def forward(self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None, inputs_embeds=None, encoder_hidden_states=None, encoder_attention_mask=None):
        outputs = ()
        for i, layer in enumerate(self.encoder.layer):
            if i in self.lora_layers:
                lora_output = self.lora_layers[f"layer_{i}"](layer.output)
                outputs = layer.output + lora_output  # Residual connection
            else:
                outputs = layer.output
        # Continue with the rest of the BERT model
        return outputs

from torch.utils.data import DataLoader, Dataset
class CustomDataset(Dataset):
    def __init__(self, data):
        self.data = data
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        return self.data[idx]
# 准备数据
dataset = CustomDataset(data)  # 假设 data 是准备好的数据集
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
# 初始化模型和优化器
model = BertModelWithLoRA(BertConfig.from_pretrained('bert-base-uncased'))
optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
# 微调
model.train()
for epoch in range(num_epochs):
    for batch in dataloader:
        inputs = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**inputs)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

model.save_pretrained('path_to_save_model')
model = BertModelWithLoRA.from_pretrained('path_to_save_model')

git clone https://github.com/microsoft/LoRA.git

git clone https://github.com/huggingface/peft.git
pip install peft