Update app.py

app.py

@@ -4,7 +4,31 @@ from knowledge_storm.lm import OpenAIModel
 from knowledge_storm.rm import YouRM
 import spaces
 import gradio as gr
+import json
+import re
 
+def convert_references_to_links(text, json_data):
+    url_mapping = json_data['url_to_unified_index']
+    
+    # Function to replace references with markdown links
+    def replace_reference(match):
+        ref_num = match.group(1)
+        url = next((url for url, index in url_mapping.items() if str(index) == ref_num), None)
+        if url:
+            return f'[{match.group(0)}]({url})'
+        return match.group(0)
+    
+    # Replace references in the text
+    processed_text = re.sub(r'\[(\d+)\]', replace_reference, text)
+    
+    # Generate reference list
+    reference_list = [f"[{index}] {url}" for url, index in sorted(url_mapping.items(), key=lambda x: x[1])]
+    
+    # Combine processed text and reference list
+    markdown_output = f"{processed_text}\n\n" + "\n".join(reference_list)
+    
+    return markdown_output
+    
 lm_configs = STORMWikiLMConfigs()
 openai_kwargs = {
     'api_key': os.getenv("OPENAI_API_KEY"),
@@ -42,7 +66,10 @@ def generate_article(prompt, progress=gr.Progress(track_tqdm=True)):
     generated_folder = prompt.replace(" ", "_")
     with open(f'outputs/{generated_folder}/storm_gen_article.txt', 'r') as file:
         content = file.read()
-    return content
+    with open(f'outputs/{generated_folder}/url_to_info.json', 'r') as file:
+        references_json = json.load(file)
+    article_full = convert_references_to_links(f'# {prompt}\n\n'+content, references_json)
+    return article_full
 
 with gr.Blocks() as demo:
     gr.Markdown("# Omnipedia article generation demo (Storm GPT-4 + You)")