From 3d30e0b9b91ed5b4d96c4e1746fe0bde1344ceec Mon Sep 17 00:00:00 2001 From: Chen Shi Date: Thu, 23 May 2024 16:56:42 +0800 Subject: [PATCH] fix: update ChatGenerationModelAdapter get strategy. (#47) --- llmuses/models/model_adapter.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/llmuses/models/model_adapter.py b/llmuses/models/model_adapter.py index ae614e39..0dfae4dd 100644 --- a/llmuses/models/model_adapter.py +++ b/llmuses/models/model_adapter.py @@ -410,8 +410,8 @@ def __init__(self, self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logger.warning(f'**Device: {self.device}') - self.generation_config = model_cfg.pop('generation_config', None) - self.generation_template = model_cfg.pop('generation_template', None) + self.generation_config = model_cfg.get('generation_config', None) + self.generation_template = model_cfg.get('generation_template', None) if self.generation_config is None or self.generation_template is None: raise ValueError('generation_config or generation_template is required for chat generation.')