chore: Update Gemini models in STORM Wiki pipeline

foreverqing · Aug 4, 2024 · 255b608 · 255b608
1 parent 34d38de
commit 255b608
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/examples/run_storm_wiki_gemini.py b/examples/run_storm_wiki_gemini.py
@@ -43,9 +43,9 @@ def main(args):
  # https://ai.google.dev/gemini-api/docs/get-started/tutorial?lang=python#list_models
  conv_simulator_lm = GoogleModel(model='models/gemini-1.5-flash', max_tokens=500, **gemini_kwargs)
  question_asker_lm = GoogleModel(model='models/gemini-1.5-flash', max_tokens=500, **gemini_kwargs)
- outline_gen_lm = GoogleModel(model='models/gemini-1.5-pro', max_tokens=400, **gemini_kwargs)
- article_gen_lm = GoogleModel(model='models/gemini-1.5-pro', max_tokens=700, **gemini_kwargs)
- article_polish_lm = GoogleModel(model='models/gemini-1.5-pro', max_tokens=4000, **gemini_kwargs)
+ outline_gen_lm = GoogleModel(model='models/gemini-1.5-pro-exp-0801', max_tokens=400, **gemini_kwargs)
+ article_gen_lm = GoogleModel(model='models/gemini-1.5-pro-exp-0801', max_tokens=700, **gemini_kwargs)
+ article_polish_lm = GoogleModel(model='models/gemini-1.5-pro-exp-0801', max_tokens=4000, **gemini_kwargs)
 
  lm_configs.set_conv_simulator_lm(conv_simulator_lm)
  lm_configs.set_question_asker_lm(question_asker_lm)