帮我优化下面的llm的隐式思维链微调代码； # 1. 时间差分方法：在训练过程中，使用了优势函数的时间差分方法，计算多步回报。# 2更复杂的奖励机制：使用了BLEU分数来评估模型生成的答案与正确答案的相似度，给予部分奖励。#3.生成通过self-play进行样本构建的代码； 
###待优化代码如下
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
import numpy as np
from tqdm import tqdm
import math
import copy
import multiprocessing
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
import os

# 使用Rich库美化输出（可选）
try:
    from rich import print
    from rich.progress import track
except ImportError:
    track = tqdm

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义思考因子（行动空间）
thought_factors = [
    '理解问题',
    '规划解题步骤',
    '执行计算',
    '检查错误',
    '调整策略',
    '总结答案',
    '验证结果',
    '反思改进'
]
num_actions = len(thought_factors)

# 蒙特卡洛树搜索节点定义
class TreeNode:
    def __init__(self, state, parent=None, action=None):
        self.state = state  # 状态文本
        self.parent = parent  # 父节点
        self.children = {}  # 子节点，键为动作索引，值为子节点
        self.action = action  # 导致当前状态的行动
        self.visits = 0  # 访问次数
        self.value = 0  # 总价值
        self.prior = 0  # 先验概率

    def is_leaf(self):
        return len(self.children) == 0

    def expand(self, action_priors):
        for action, prior in enumerate(action_priors):
            if action not in self.children:
                new_state = self.state + '\n' + thought_factors[action] + ':'
                self.children[action] = TreeNode(new_state, parent=self, action=action)
                self.children[action].prior = prior

    def backpropagate(self, value):
        self.visits += 1
        self.value += value
        if self.parent:
            self.parent.backpropagate(value)

    def select(self, c_puct):
        # 使用并行化时，需要注意线程安全问题
        best_score = -float('inf')
        best_action = None
        best_child = None
        for action, child in self.children.items():
            ucb_score = (child.value / (child.visits + 1e-5)) + \
                        c_puct * child.prior * math.sqrt(self.visits) / (1 + child.visits)
            if ucb_score > best_score:
                best_score = ucb_score
                best_action = action
                best_child = child
        return best_action, best_child

# 模型定义
class LLMActorCritic(nn.Module):
    def __init__(self, model_name='Qwen/Qwen-7B'):
        super(LLMActorCritic, self).__init__()
        self.llm = AutoModelForCausalLM.from_pretrained(model_name)
        self.hidden_size = self.llm.config.hidden_size

        # 策略网络（Actor）
        self.policy_head = nn.Linear(self.hidden_size, num_actions)

        # 价值网络（Critic）
        self.value_head = nn.Linear(self.hidden_size, 1)

    def forward(self, input_ids, attention_mask=None):
        outputs = self.llm.base_model(input_ids=input_ids, attention_mask=attention_mask)
        last_hidden_state = outputs.last_hidden_state  # (batch_size, seq_length, hidden_size)

        # 使用序列最后一个Token的表示
        hidden_state = last_hidden_state[:, -1, :]  # (batch_size, hidden_size)

        # 策略网络输出
        policy_logits = self.policy_head(hidden_state)  # (batch_size, num_actions)
        policy_probs = F.softmax(policy_logits, dim=-1)

        # 价值网络输出
        value = self.value_head(hidden_state)  # (batch_size, 1)

        return policy_probs, value

    def generate_tokens(self, input_ids, max_length=50):
        # 使用LLM生成后续的Token序列
        output_ids = self.llm.generate(
            input_ids.to(device),
            max_length=max_length,
            do_sample=True,
            top_p=0.95,
            top_k=50,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id
        )
        return output_ids

# 蒙特卡洛树搜索实现
class MCTS:
    def __init__(self, model, tokenizer, max_depth=5, c_puct=1.0, num_simulations=50):
        self.model = model
        self.tokenizer = tokenizer
        self.max_depth = max_depth
        self.c_puct = c_puct
        self.num_simulations = num_simulations

    def simulate(self, node):
        # 如果到达最大深度，返回估计价值
        if node.is_leaf() or node.visits >= self.max_depth:
            input_ids = self.tokenizer.encode(node.state, return_tensors='pt', truncation=True).to(device)
            attention_mask = torch.ones_like(input_ids).to(device)
            with torch.no_grad():
                policy_probs, value = self.model(input_ids, attention_mask)
            node.expand(policy_probs.cpu().numpy()[0])
            return value.item()
        else:
            action, child = node.select(self.c_puct)
            value = self.simulate(child)
            node.backpropagate(value)
            return value

    def search(self, root):
        # 并行模拟
        with ThreadPoolExecutor(max_workers=4) as executor:
            futures = [executor.submit(self.simulate, root) for _ in range(self.num_simulations)]
            for _ in futures:
                pass

        # 从根节点选择访问次数最多的行动
        visits = {action: child.visits for action, child in root.children.items()}
        best_action = max(visits, key=visits.get)
        return root.children[best_action]

# 奖励评估函数
def evaluate_answer(generated_answer, correct_answer):
    # 计算BLEU分数作为部分奖励
    from nltk.translate.bleu_score import sentence_bleu
    reference = [correct_answer.split()]
    candidate = generated_answer.split()
    bleu_score = sentence_bleu(reference, candidate)
    # 将BLEU分数缩放到[-1, 1]范围
    reward = 2 * bleu_score - 1
    return reward

# 训练函数
def train(model, data, tokenizer, num_epochs=5):
    model.train()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)

    for epoch in range(num_epochs):
        epoch_loss = 0
        for sample in track(data, description=f"Epoch {epoch+1}"):
            question = sample['question']
            correct_answer = sample['correct_answer']

            # 初始状态
            initial_state = question
            root = TreeNode(initial_state)
            mcts = MCTS(model, tokenizer)

            # 收集序列
            states = []
            actions = []
            rewards = []
            values = []

            # 思维过程
            generated_cot = ''
            node = root

            for step in range(5):  # 最大步骤
                # MCTS搜索
                best_child = mcts.search(node)
                action_idx = best_child.action
                thought_factor = thought_factors[action_idx]
                actions.append(action_idx)
                states.append(node.state)

                # 生成思维片段
                cot_input = best_child.state
                cot_input_ids = tokenizer.encode(cot_input, return_tensors='pt', truncation=True).to(device)
                with torch.no_grad():
                    output_ids = model.generate_tokens(cot_input_ids)
                output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
                new_cot = output_text[len(tokenizer.decode(cot_input_ids[0], skip_special_tokens=True)):].strip()
                generated_cot = cot_input + new_cot

                # 更新节点
                node = best_child

                # 判断是否结束
                if '答案' in new_cot or '解答' in new_cot:
                    break

            # 评估最终答案
            generated_answer = extract_answer_from_cot(generated_cot)
            reward = evaluate_answer(generated_answer, correct_answer)

            # 奖励传播
            rewards = [0] * (len(actions) - 1) + [reward]

            # 计算优势函数和时间差分目标
            G = 0
            gammas = 0.99
            returns = []
            for r in reversed(rewards):
                G = r + gammas * G
                returns.insert(0, G)

            # 更新策略和价值网络
            optimizer.zero_grad()

            policy_losses = []
            value_losses = []
            entropies = []

            for t in range(len(actions)):
                state = states[t]
                action = actions[t]
                R = returns[t]

                input_ids = tokenizer.encode(state, return_tensors='pt', truncation=True).to(device)
                attention_mask = torch.ones_like(input_ids).to(device)

                policy_probs, value = model(input_ids, attention_mask)
                log_prob = torch.log(policy_probs.squeeze(0)[action] + 1e-5)
                advantage = R - value.item()

                policy_losses.append(-log_prob * advantage)
                value_losses.append(F.mse_loss(value.squeeze(), torch.tensor([R], device=device)))
                entropies.append(- (policy_probs * torch.log(policy_probs + 1e-5)).sum())

            loss = torch.stack(policy_losses).sum() + torch.stack(value_losses).sum() - 0.01 * torch.stack(entropies).sum()

            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()

        print(f"Epoch {epoch+1}, Loss: {epoch_loss/len(data)}")

        # 定期保存模型
        model_save_path = f"llm_actor_critic_epoch{epoch+1}.pt"
        torch.save(model.state_dict(), model_save_path)
        print(f"Model saved to {model_save_path}")

# 提取答案的函数
def extract_answer_from_cot(generated_cot):
    lines = generated_cot.strip().split('\n')
    for line in reversed(lines):
        if '答案' in line or '解答' in line:
            return line.split(':')[-1].strip()
    return '未找到答案'

# 加载模型和Tokenizer
model_name = 'Qwen/Qwen-7B'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = LLMActorCritic(model_name).to(device)

# 示例数据
data = [
    {
        'question': '两个数的和是10，乘积是21，求这两个数。',
        'correct_answer': '7和3'
    },
    # 可以添加更多数据
]

# 开始训练
train(model, data, tokenizer, num_epochs=5)

# 推理
def generate_answer(model, tokenizer, question, max_steps=5):
    model.eval()
    with torch.no_grad():
        initial_state = question
        root = TreeNode(initial_state)
        mcts = MCTS(model, tokenizer)

        generated_cot = ''
        node = root

        for step in range(max_steps):
            # 使用MCTS选择最佳行动
            best_child = mcts.search(node)
            action_idx = best_child.action
            thought_factor = thought_factors[action_idx]

            # 根据思考因子生成对应的文本片段
            cot_input = best_child.state
            cot_input_ids = tokenizer.encode(cot_input, return_tensors='pt', truncation=True).to(device)
            output_ids = model.generate_tokens(cot_input_ids)