fix: LLMTrainer supports lr param

llm_chat_improved.py

@@ -580,11 +580,11 @@ class ConversationDataset(Dataset):
 class LLMTrainer:
     """LLM训练器"""
     
-    def __init__(self, model: nn.Module, device: str = 'cpu', use_amp: bool = None):
+    def __init__(self, model: nn.Module, device: str = 'cpu', lr: float = 0.0001, use_amp: bool = None):
         self.model = model.to(device)
         self.device = device
         self.criterion = nn.CrossEntropyLoss(ignore_index=0)
-        self.optimizer = optim.AdamW(model.parameters(), lr=0.0001, weight_decay=0.01)
+        self.optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
         self.scheduler = optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=100)
         
         # 混合精度训练