diff --git a/pilot/model/adapter.py b/pilot/model/adapter.py
index 900d51d4a..05643c734 100644
--- a/pilot/model/adapter.py
+++ b/pilot/model/adapter.py
@@ -12,6 +12,8 @@ from transformers import (
     LlamaTokenizer,
     BitsAndBytesConfig,
 )
+from transformers.generation.utils import GenerationConfig
+
 from pilot.configs.model_config import DEVICE
 from pilot.configs.config import Config
 
diff --git a/pilot/model/conversation.py b/pilot/model/conversation.py
index 11ca03ed8..d29e73da3 100644
--- a/pilot/model/conversation.py
+++ b/pilot/model/conversation.py
@@ -2,6 +2,8 @@
 Fork from fastchat: https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
 
 Conversation prompt templates.
+
+
 """
 
 import dataclasses