Update app.py

app.py

@@ -2,6 +2,7 @@ import os
 from knowledge_storm import STORMWikiRunnerArguments, STORMWikiRunner, STORMWikiLMConfigs
 from knowledge_storm.lm import OpenAIModel
 from knowledge_storm.rm import YouRM
+import spaces
 
 lm_configs = STORMWikiLMConfigs()
 openai_kwargs = {
@@ -25,6 +26,7 @@ engine_args = STORMWikiRunnerArguments("outputs")
 rm = YouRM(ydc_api_key=os.getenv('YDC_API_KEY'), k=engine_args.search_top_k)
 runner = STORMWikiRunner(engine_args, lm_configs, rm)
 
+@spaces.GPU
 def generate_article(prompt, progress=gr.Progress(track_tqdm=True)):
     runner.run(
         topic=prompt,
@@ -36,4 +38,11 @@ def generate_article(prompt, progress=gr.Progress(track_tqdm=True)):
     runner.post_run()
     runner.summary()
 
+with gr.Blocks() as demo:
+    prompt = gr.Textbox(label="Prompt")
+    output = gr.Markdown(label="Output")
+    btn = gr.Button("Generate")
+    btn.click(fn=generate_article, inputs=prompt, outputs=output)
 
+if __name__ == "__main__":
+    demo.launch()