(args)
| 30 | |
| 31 | |
| 32 | def main(args): |
| 33 | load_api_key(toml_file_path="secrets.toml") |
| 34 | lm_configs = STORMWikiLMConfigs() |
| 35 | claude_kwargs = { |
| 36 | "api_key": os.getenv("ANTHROPIC_API_KEY"), |
| 37 | "temperature": 1.0, |
| 38 | "top_p": 0.9, |
| 39 | } |
| 40 | |
| 41 | # STORM is a LM system so different components can be powered by different models. |
| 42 | # For a good balance between cost and quality, you can choose a cheaper/faster model for conv_simulator_lm |
| 43 | # which is used to split queries, synthesize answers in the conversation. We recommend using stronger models |
| 44 | # for outline_gen_lm which is responsible for organizing the collected information, and article_gen_lm |
| 45 | # which is responsible for generating sections with citations. |
| 46 | conv_simulator_lm = ClaudeModel( |
| 47 | model="claude-3-haiku-20240307", max_tokens=500, **claude_kwargs |
| 48 | ) |
| 49 | question_asker_lm = ClaudeModel( |
| 50 | model="claude-3-sonnet-20240229", max_tokens=500, **claude_kwargs |
| 51 | ) |
| 52 | outline_gen_lm = ClaudeModel( |
| 53 | model="claude-3-opus-20240229", max_tokens=400, **claude_kwargs |
| 54 | ) |
| 55 | article_gen_lm = ClaudeModel( |
| 56 | model="claude-3-opus-20240229", max_tokens=700, **claude_kwargs |
| 57 | ) |
| 58 | article_polish_lm = ClaudeModel( |
| 59 | model="claude-3-opus-20240229", max_tokens=4000, **claude_kwargs |
| 60 | ) |
| 61 | |
| 62 | lm_configs.set_conv_simulator_lm(conv_simulator_lm) |
| 63 | lm_configs.set_question_asker_lm(question_asker_lm) |
| 64 | lm_configs.set_outline_gen_lm(outline_gen_lm) |
| 65 | lm_configs.set_article_gen_lm(article_gen_lm) |
| 66 | lm_configs.set_article_polish_lm(article_polish_lm) |
| 67 | |
| 68 | engine_args = STORMWikiRunnerArguments( |
| 69 | output_dir=args.output_dir, |
| 70 | max_conv_turn=args.max_conv_turn, |
| 71 | max_perspective=args.max_perspective, |
| 72 | search_top_k=args.search_top_k, |
| 73 | max_thread_num=args.max_thread_num, |
| 74 | ) |
| 75 | # Documentation to generate the data is available here: |
| 76 | # https://serper.dev/playground |
| 77 | # Important to note that tbs(date range is hardcoded values). |
| 78 | # num is results per pages and is recommended to use in increments of 10(10, 20, etc). |
| 79 | # page is how many pages will be searched. |
| 80 | # h1 is where the google search will orginate from. |
| 81 | topic = input("topic: ") |
| 82 | data = {"autocorrect": True, "num": 10, "page": 1} |
| 83 | rm = SerperRM(serper_search_api_key=os.getenv("SERPER_API_KEY"), query_params=data) |
| 84 | |
| 85 | runner = STORMWikiRunner(engine_args, lm_configs, rm) |
| 86 | |
| 87 | runner.run( |
| 88 | topic=topic, |
| 89 | do_research=args.do_research, |
no test coverage detected