文字语义理解技术中的语义角色标注问题

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

class SRLDataset(Dataset):
    def __init__(self, sentences, labels):
        self.sentences = sentences
        self.labels = labels
        
    def __len__(self):
        return len(self.sentences)
    
    def __getitem__(self, idx):
        sentence = self.sentences[idx]
        label = self.labels[idx]
        return sentence, label

class BiLSTMCRF(nn.Module):
    def __init__(self, embedding_dim, hidden_dim, num_classes):
        super(BiLSTMCRF, self).__init__()
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.num_classes = num_classes
        
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim // 2, bidirectional=True)
        self.hidden2tag = nn.Linear(hidden_dim, num_classes)
        self.crf = CRF(num_classes)
        
    def forward(self, sentence):
        embeds = self.embedding(sentence)
        lstm_out, _ = self.lstm(embeds)
        tag_space = self.hidden2tag(lstm_out)
        return tag_space
    
    def loss(self, sentence, targets):
        forward_score = self.forward(sentence)
        return self.crf.loss(forward_score, targets)
        
    def decode(self, sentence):
        forward_score = self.forward(sentence)
        return self.crf.decode(forward_score)

# 数据准备
sentences = [['小明', '吃了', '一个', '苹果'], ['小明', '吃了', '一个', '香蕉']]
labels = [['施事者', '动作', '数量', '受事者'], ['施事者', '动作', '数量', '受事者']]
dataset = SRLDataset(sentences, labels)

# 模型训练
model = BiLSTMCRF(embedding_dim, hidden_dim, num_classes)
optimizer = optim.SGD(model.parameters(), lr=0.1)
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

for epoch in range(epochs):
    for sentence, targets in data_loader:
        optimizer.zero_grad()
        sentence = torch.tensor(sentence)
        targets = torch.tensor(targets)
        loss = model.loss(sentence, targets)
        loss.backward()
        optimizer.step()

# 模型预测
new_sentence = [['小明', '去了', '中国']]
new_sentence = torch.tensor(new_sentence)
predicted_labels = model.decode(new_sentence)
print(predicted_labels)