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

torch.manual_seed(1)

START_TAG = "<START>"
STOP_TAG = "<STOP>"
EMBEDDING_DIM = 50
HIDDEN_DIM = 50


def read_data(filepath):
    sentences = []
    tags = []
    with open(filepath, 'r', encoding='utf-8') as f:
        sentence = []
        tag_seq = []
        for line in f:
            if line.strip() == "":
                if sentence:
                    sentences.append(sentence)
                    tags.append(tag_seq)
                    sentence = []
                    tag_seq = []
            else:
                word, tag = line.strip().split()
                sentence.append(word)
                tag_seq.append(tag)
        if sentence:
            sentences.append(sentence)
            tags.append(tag_seq)
    return sentences, tags


class NerDataset(Dataset):
    def __init__(self, filepath, word_to_ix, tag_to_ix):
        self.sentences, self.tags = read_data(filepath)
        self.word_to_ix = word_to_ix
        self.tag_to_ix = tag_to_ix

    def __len__(self):
        return len(self.sentences)

    def __getitem__(self, idx):
        sentence = self.sentences[idx]
        tags = self.tags[idx]
        return ([self.word_to_ix.get(word, self.word_to_ix["<UNK>"]) for word in sentence],
                [self.tag_to_ix[tag] for tag in tags])


def argmax(vec):
    _, idx = torch.max(vec, 1)
    return idx.item()


def prepare_sequence(seq, to_ix):
    idxs = [to_ix[w] for w in seq]
    return torch.tensor(idxs, dtype=torch.long)


def log_sum_exp(vec):
    max_score = vec[0, argmax(vec)]
    max_score_broadcast = max_score.view(1, -1).expand(1, vec.size()[1])
    return max_score + torch.log(torch.sum(torch.exp(vec - max_score_broadcast)))


class BiLSTM_CRF(nn.Module):
    def __init__(self, vocab_size, tag_to_ix, embedding_dim, hidden_dim):
        super(BiLSTM_CRF, self).__init__()
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.vocab_size = vocab_size
        self.tag_to_ix = tag_to_ix
        self.tagset_size = len(tag_to_ix)
        self.word_embeds = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim // 2, num_layers=1, bidirectional=True)
        self.hidden2tag = nn.Linear(hidden_dim, self.tagset_size)
        self.transitions = nn.Parameter(torch.randn(self.tagset_size, self.tagset_size))
        self.transitions.data[tag_to_ix[START_TAG], :] = -10000
        self.transitions.data[:, tag_to_ix[STOP_TAG]] = -10000

    def _forward_alg(self, feats):
        init_alphas = torch.full((1, self.tagset_size), -10000.)
        init_alphas[0][self.tag_to_ix[START_TAG]] = 0.
        forward_var = init_alphas
        for feat in feats:
            alphas_t = []
            for next_tag in range(self.tagset_size):
                emit_score = feat[next_tag].view(1, -1).expand(1, self.tagset_size)
                trans_score = self.transitions[next_tag].view(1, -1)
                next_tag_var = forward_var + trans_score + emit_score
                alphas_t.append(log_sum_exp(next_tag_var).view(1))
            forward_var = torch.cat(alphas_t).view(1, -1)
        terminal_var = forward_var + self.transitions[self.tag_to_ix[STOP_TAG]]
        alpha = log_sum_exp(terminal_var)
        return alpha

    def _get_lstm_features(self, sentence):
        self.hidden = self.init_hidden()
        embeds = self.word_embeds(sentence).view(len(sentence), 1, -1)
        lstm_out, self.hidden = self.lstm(embeds)
        lstm_feats = self.hidden2tag(lstm_out.view(len(sentence), -1))
        return lstm_feats

    def _score_sentence(self, feats, tags):
        score = torch.zeros(1)
        tags = torch.cat([torch.tensor([self.tag_to_ix[START_TAG]], dtype=torch.long), tags])
        for i, feat in enumerate(feats):
            score = score + self.transitions[tags[i + 1], tags[i]] + feat[tags[i + 1]]
        score = score + self.transitions[self.tag_to_ix[STOP_TAG], tags[-1]]
        return score

    def _viterbi_decode(self, features):
        backpointers = []
        # 初始化动态规划矩阵的起始值，所有位置设为一个很小的数，除了开始标签的位置
        init_scores = torch.full((1, self.tagset_size), -10000.)
        init_scores[0][self.tag_to_ix[START_TAG]] = 0

        # 初始化前向变量为初始分数
        forward_scores = init_scores
        for feature in features:
            step_backpointers = []
            step_scores = []

            for next_tag in range(self.tagset_size):
                # 计算到达下一个标签的分数
                next_tag_scores = forward_scores + self.transitions[next_tag]
                best_tag_id = argmax(next_tag_scores)
                step_backpointers.append(best_tag_id)
                step_scores.append(next_tag_scores[0][best_tag_id].view(1))
            # 更新前向分数为当前步骤的分数加上特征分数
            forward_scores = (torch.cat(step_scores) + feature).view(1, -1)
            backpointers.append(step_backpointers)

        terminal_scores = forward_scores + self.transitions[self.tag_to_ix[STOP_TAG]]
        best_tag_id = argmax(terminal_scores)
        final_path_score = terminal_scores[0][best_tag_id]

        # 回溯最佳路径
        best_path = [best_tag_id]
        for step_backpointers in reversed(backpointers):
            best_tag_id = step_backpointers[best_tag_id]
            best_path.append(best_tag_id)
        # 检查路径的起始是否正确
        start = best_path.pop()
        assert start == self.tag_to_ix[START_TAG]  # 确认起点是开始标签
        best_path.reverse()  # 翻转路径得到正确的顺序
        return final_path_score, best_path

    def forward(self, sentence):
        lstm_features = self._get_lstm_features(sentence)
        path_score, best_tag_sequence = self._viterbi_decode(lstm_features)
        return path_score, best_tag_sequence

    def neg_log_likelihood(self, sentence, tags):
        feats = self._get_lstm_features(sentence)
        forward_score = self._forward_alg(feats)
        gold_score = self._score_sentence(feats, tags)
        return forward_score - gold_score

    def init_hidden(self):
        return torch.randn(2, 1, self.hidden_dim // 2), torch.randn(2, 1, self.hidden_dim // 2)


# 数据预处理
def build_vocab(data):
    word_to_ix = {"<UNK>": 0}
    tag_to_ix = {START_TAG: 0, STOP_TAG: 1}
    for sentence, tags in data:
        for word in sentence:
            if word not in word_to_ix:
                word_to_ix[word] = len(word_to_ix)
        for tag in tags:
            if tag not in tag_to_ix:
                tag_to_ix[tag] = len(tag_to_ix)
    return word_to_ix, tag_to_ix


# 读取数据
sentences, tags = read_data('train.txt')
train_data = list(zip(sentences, tags))

# 构建词典
word_to_ix, tag_to_ix = build_vocab(train_data)

# 创建数据集
train_dataset = NerDataset('train.txt', word_to_ix, tag_to_ix)
val_dataset = NerDataset('validation.txt', word_to_ix, tag_to_ix)

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False)

# 创建模型
model = BiLSTM_CRF(len(word_to_ix), tag_to_ix, EMBEDDING_DIM, HIDDEN_DIM)
optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay=1e-4)

# 训练模型
if os.path.exists('model_complete.pth'):
    model = torch.load('model_complete.pth')
for epoch in range(1):
    for sentence, tags in train_loader:
        model.zero_grad()
        sentence_in = sentence.clone().detach()
        targets = tags.clone().detach()
        loss = model.neg_log_likelihood(sentence_in, targets)
        loss.backward()
        optimizer.step()
torch.save(model, 'model_complete.pth')


def evaluate_and_write_results(model, validation_loader, ix_to_tag, raw_sentences, file_path):
    model.eval()
    with torch.no_grad():
        with open(file_path, 'w', encoding='utf-8') as f:
            for i, (sentence, tags) in enumerate(validation_loader):
                raw_sentence = raw_sentences[i]
                for j in range(len(sentence)):
                    sentence_in = sentence[j].clone().detach()
                    _, predicted_tag_indice = model(sentence_in)
                    predicted_tag = ix_to_tag[predicted_tag_indice]
                    f.write(f"{raw_sentence[j]} {predicted_tag}\n")
                f.write("\n")
    print(f"解码结果已写入 {file_path} 文件中。")


ix_to_tag = {ix: tag for tag, ix in tag_to_ix.items()}
raw_sentences, tags = read_data('validation.txt')
evaluate_and_write_results(model, val_loader, ix_to_tag, raw_sentences, "BiLSTM_CRF_model.txt")

我的代码哪里有问题