File tree Expand file tree Collapse file tree 1 file changed +9
-6
lines changed Expand file tree Collapse file tree 1 file changed +9
-6
lines changed Original file line number Diff line number Diff line change @@ -152,12 +152,15 @@ def _create_llm(self, llm_config: dict) -> object:
152152 raise ValueError (f"""Provider { llm_params ['model_provider' ]} is not supported.
153153 If possible, try to use a model instance instead.""" )
154154
155- try :
156- self .model_token = models_tokens [llm_params ["model_provider" ]][llm_params ["model" ]]
157- except KeyError :
158- print (f"""Model { llm_params ['model_provider' ]} /{ llm_params ['model' ]} not found,
159- using default token size (8192)""" )
160- self .model_token = 8192
155+ if "model_tokens" not in llm_params :
156+ try :
157+ self .model_token = models_tokens [llm_params ["model_provider" ]][llm_params ["model" ]]
158+ except KeyError :
159+ print (f"""Model { llm_params ['model_provider' ]} /{ llm_params ['model' ]} not found,
160+ using default token size (8192)""" )
161+ self .model_token = 8192
162+ else :
163+ self .model_token = llm_params ["model_tokens" ]
161164
162165 try :
163166 if llm_params ["model_provider" ] not in \
You can’t perform that action at this time.
0 commit comments