File size: 5,940 Bytes
80a598c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
import time
import json
import pandas as pd
import streamlit as st
from dotenv import load_dotenv
from http import HTTPStatus
from src.lm import QwenModel
from src.rm import GoogleSearchAli_new
import sys
sys.path.append('./src/DeepThink/modules')
from mindmap import MindMap
from storm_dataclass import Article
from article_generation import ArticleGenerationModule
from article_polish import ArticlePolishingModule
from outline_generation import OutlineGenerationModule
import os
import subprocess
bash_command = "pip install --upgrade pip"
process = subprocess.Popen(bash_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
# Load environment variables and API keys
# load_dotenv()
openai_kwargs = {
'api_key': os.getenv("OPENAI_API_KEY"),
'api_provider': os.getenv('OPENAI_API_TYPE'),
'temperature': 1.0,
'top_p': 0.9,
'api_base': os.getenv('AZURE_API_BASE'),
'api_version': os.getenv('AZURE_API_VERSION'),
}
lm = QwenModel(model='qwen-plus', max_tokens=1000, **openai_kwargs)
lm4outline = QwenModel(model='qwen-plus', max_tokens=1000, **openai_kwargs)
lm4gensection = QwenModel(model='qwen-plus', max_tokens=2000, **openai_kwargs)
lm4polish = QwenModel(model='qwen-plus', max_tokens=4000, **openai_kwargs)
rm = GoogleSearchAli_new(k=5)
st.set_page_config(page_title='OmniThink', layout="wide")
st.warning("Announcement: Due to the recent high volume of visitors, search API quota limitations, you may encounter an error: "
"'ValueError: Expected 2D array, got 1D array instead: array=[]. "
"Reshape your data either using array.reshape(-1, 1) if your data has a single feature "
"or array.reshape(1, -1) if it contains a single sample.' "
"If this error occurs, please try again in a few hours.")
st.title('🤔 OmniThink')
st.markdown('_OmniThink is a tool that helps you think deeply about a topic, generate an outline, and write an article._')
# Sidebar for configuration and examples
with st.sidebar:
st.header('Configuration')
MAX_ROUNDS = st.number_input('Retrieval Depth', min_value=0, max_value=10, value=2, step=1)
models = ['Qwen-Plus', 'Coming Soon']
selected_example = st.selectbox('LLM:', models)
searchers = ['GoogleSearch', 'Coming Soon']
selected_example = st.selectbox('Search engine', searchers)
n_max_doc = st.number_input('Number of web pages retrievad in single search', min_value=1, max_value=50, value=10, step=5)
st.header('Examples')
examples = ['AlphaFold', '2024 Hualien City Earthquake', 'Taylor Swift', 'Yoon Seok-youl']
selected_example = st.selectbox('case', examples)
status_placeholder = st.empty()
mind_map = MindMap(
retriever=rm,
gen_concept_lm = lm4outline,
gen_concept_lm2 = lm4outline,
search_top_k = n_max_doc,
depth= MAX_ROUNDS
)
def Think(input_topic):
generator = mind_map.build_map(input_topic)
st.markdown(f'Performing an in-depth search on the content related to {input_topic}...')
for idx, layer in enumerate(generator):
print(layer)
print('layer!!!')
st.markdown(f'Deep Thinking Retrieval at Level {idx + 1}...')
status_placeholder.text(f"Currently conducting the {idx + 1}th level deep thinking retrieval, estimated to take {(idx+1)*3} minutes.")
for node in layer:
category = node.category
print(f'category: {category}')
with st.expander(f'{category}'):
st.markdown(f'### The concept of {node.category}')
print(node.concept)
for concept in node.concept:
st.markdown(f'* {concept}')
st.markdown(f'### The web of {node.category}')
for idx, info in enumerate(node.info):
st.markdown(f'{idx + 1}. {info["title"]} \n {info["snippets"]}')
st.markdown(f'Constructing an index table for the {mind_map.get_web_number()} retrieved web pages...')
mind_map.prepare_table_for_retrieval()
return '__finish__', '__finish__'
def GenOutline(input_topic):
status_placeholder.text("The outline writing is in progress and is expected to take 1 minute.")
ogm = OutlineGenerationModule(lm)
outline = ogm.generate_outline(topic= input_topic, mindmap = mind_map)
return outline
def GenArticle(input_topic, outline):
status_placeholder.text("The article writing is in progress and is expected to take 3 minutes.")
article_with_outline = Article.from_outline_str(topic=input_topic, outline_str=outline)
ag = ArticleGenerationModule(retriever = rm, article_gen_lm = lm, retrieve_top_k = 3, max_thread_num = 10)
article = ag.generate_article(topic = topic, mindmap = mind_map, article_with_outline = article_with_outline)
ap = ArticlePolishingModule(article_gen_lm = lm, article_polish_lm = lm)
article = ap.polish_article(topic = topic, draft_article = article)
return article.to_string()
with st.form('my_form'):
topic = st.text_input('Please enter the topic you are interested in.', value=selected_example, placeholder='Please enter the topic you are interested in.')
submit_button = st.form_submit_button('Generate！')
if submit_button:
if topic:
st.markdown('### Thought process')
summary, news_timeline = Think(topic)
st.session_state.summary = summary
st.session_state.news_timeline = news_timeline
st.markdown('### Outline generation')
with st.expander("Outline generation", expanded=True):
outline = GenOutline(topic)
st.text(outline)
st.markdown('### article generation')
with st.expander("article generation", expanded=True):
article = GenArticle(topic, outline)
st.markdown(article)
else:
st.error('Please enter the subject.')
|