ShawnAI/VectorDB-ChatBot
4
1import gradio as gr2import random3import time4 5from langchain import PromptTemplate6from langchain.llms import OpenAI7from langchain.chat_models import ChatOpenAI8from langchain.embeddings import HuggingFaceEmbeddings, HuggingFaceInstructEmbeddings, OpenAIEmbeddings9from langchain.vectorstores import Pinecone10from langchain.chains import LLMChain11from langchain.chains.question_answering import load_qa_chain12import pinecone13 14import os15os.environ["TOKENIZERS_PARALLELISM"] = "false"16 17#OPENAI_API_KEY = ""18OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")19OPENAI_TEMP = 120OPENAI_API_LINK = "[OpenAI API Key](https://platform.openai.com/account/api-keys)"21OPENAI_LINK = "[OpenAI](https://openai.com)"22 23PINECONE_KEY = os.environ.get("PINECONE_KEY", "")24PINECONE_ENV = os.environ.get("PINECONE_ENV", "asia-northeast1-gcp")25PINECONE_INDEX = os.environ.get("PINECONE_INDEX", '3gpp-r16')26 27PINECONE_LINK = "[Pinecone](https://www.pinecone.io)"28LANGCHAIN_LINK = "[LangChain](https://python.langchain.com/en/latest/index.html)"29 30EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "hkunlp/instructor-large")31EMBEDDING_LOADER = os.environ.get("EMBEDDING_LOADER", "HuggingFaceInstructEmbeddings")32EMBEDDING_LIST = ["HuggingFaceInstructEmbeddings", "HuggingFaceEmbeddings", "OpenAIEmbeddings"]33 34# return top-k text chunks from vector store35TOP_K_DEFAULT = 1536TOP_K_MAX = 3037SCORE_DEFAULT = 0.3338 39 40BUTTON_MIN_WIDTH = 21541 42LLM_NULL = "LLM-UNLOAD-critical"43LLM_DONE = "LLM-LOADED-9cf"44 45DB_NULL = "DB-UNLOAD-critical"46DB_DONE = "DB-LOADED-9cf"47 48FORK_BADGE = "Fork-HuggingFace Space-9cf"49 50 51def get_logo(inputs, logo) -> str:52 return f"""https://img.shields.io/badge/{inputs}?style=flat&logo={logo}&logoColor=white"""53 54def get_status(inputs, logo, pos) -> str:55 return f"""<img56 src = "{get_logo(inputs, logo)}";57 style = "margin: 0 auto;float:{pos};border: 2px solid transparent;";58 >"""59 60 61KEY_INIT = "Initialize Model"62KEY_SUBMIT = "Submit"63KEY_CLEAR = "Clear"64 65MODEL_NULL = get_status(LLM_NULL, "openai", "right")66MODEL_DONE = get_status(LLM_DONE, "openai", "right")67 68DOCS_NULL = get_status(DB_NULL, "processingfoundation", "right")69DOCS_DONE = get_status(DB_DONE, "processingfoundation", "right")70 71TAB_1 = "Chatbot"72TAB_2 = "Details"73TAB_3 = "Database"74TAB_4 = "TODO"75 76 77 78FAVICON = './icon.svg'79 80LLM_LIST = ["gpt-3.5-turbo", "text-davinci-003"]81 82 83DOC_1 = '3GPP'84DOC_2 = 'HTTP2'85 86DOC_SUPPORTED = [DOC_1]87DOC_DEFAULT = [DOC_1]88DOC_LABEL = "Reference Docs"89 90 91MODEL_WARNING = f"Please paste your **{OPENAI_API_LINK}** and then **{KEY_INIT}**"92 93DOCS_WARNING = f"""Database Unloaded94Please check your **{TAB_3}** config and then **{KEY_INIT}**95Or you could uncheck **{DOC_LABEL}** to ask LLM directly"""96 97 98webui_title = """99# OpenAI Chatbot Based on Vector Database100"""101 102dup_link = f'''<a href="https://huggingface.co/spaces/ShawnAI/VectorDB-ChatBot?duplicate=true"103style="display:grid; width: 200px;">104<img src="{get_logo(FORK_BADGE, "addthis")}"></a>'''105 106init_message = f"""This demonstration website is based on \107**{OPENAI_LINK}** with **{LANGCHAIN_LINK}** and **{PINECONE_LINK}**108 1. Insert your **{OPENAI_API_LINK}** and click `{KEY_INIT}`109 2. Insert your **Question** and click `{KEY_SUBMIT}`110"""111 112PROMPT_DOC = PromptTemplate(113 input_variables=["context", "chat_history", "question"],114 template="""Context:115##116{context}117##118 119Chat History:120##121{chat_history}122##123 124Question:125{question}126 127Answer:"""128)129 130PROMPT_BASE = PromptTemplate(131 input_variables=['question', "chat_history"],132 template="""Chat History:133##134{chat_history}135##136 137Question:138##139{question}140##141 142Answer:"""143)144 145#----------------------------------------------------------------------------------------------------------146#----------------------------------------------------------------------------------------------------------147def init_rwkv():148 try:149 import rwkv150 return True151 except Exception:152 print("RWKV not found, skip local llm")153 return False154 155 156def init_model(api_key, emb_name, emb_loader, db_api_key, db_env, db_index):157 init_rwkv()158 try:159 if not (api_key and api_key.startswith("sk-") and len(api_key) > 50):160 return None,MODEL_NULL+DOCS_NULL,None,None,None,None161 162 163 164 llm_dict = {}165 for llm_name in LLM_LIST:166 if llm_name == "gpt-3.5-turbo":167 llm_dict[llm_name] = ChatOpenAI(model_name=llm_name,168 temperature = OPENAI_TEMP,169 openai_api_key = api_key170 )171 else:172 llm_dict[llm_name] = OpenAI(model_name=llm_name,173 temperature = OPENAI_TEMP,174 openai_api_key = api_key)175 176 if not (emb_name and db_api_key and db_env and db_index):177 return api_key,MODEL_DONE+DOCS_NULL,llm_dict,None,None,None178 179 if emb_loader == "OpenAIEmbeddings":180 embeddings = eval(emb_loader)(openai_api_key=api_key)181 else:182 embeddings = eval(emb_loader)(model_name=emb_name)183 184 pinecone.init(api_key = db_api_key,185 environment = db_env)186 db = Pinecone.from_existing_index(index_name = db_index,187 embedding = embeddings)188 189 return api_key, MODEL_DONE+DOCS_DONE, llm_dict, None, db, None190 191 except Exception as e:192 print(e)193 return None,MODEL_NULL+DOCS_NULL,None,None,None,None194 195 196def get_chat_history(inputs) -> str:197 res = []198 for human, ai in inputs:199 res.append(f"Q: {human}\nA: {ai}")200 return "\n".join(res)201 202def remove_duplicates(documents, score_min):203 seen_content = set()204 unique_documents = []205 for (doc, score) in documents:206 if (doc.page_content not in seen_content) and (score >= score_min):207 seen_content.add(doc.page_content)208 unique_documents.append(doc)209 return unique_documents210 211def doc_similarity(query, db, top_k, score):212 docs = db.similarity_search_with_score(query = query,213 k=top_k)214 #docsearch = db.as_retriever(search_kwargs={'k':top_k})215 #docs = docsearch.get_relevant_documents(query)216 udocs = remove_duplicates(docs, score)217 return udocs218 219def user(user_message, history):220 return "", history+[[user_message, None]]221 222def bot(box_message, ref_message,223 llm_dropdown, llm_dict, doc_list,224 db, top_k, score):225 226 # bot_message = random.choice(["Yes", "No"])227 # 0 is user question, 1 is bot response228 question = box_message[-1][0]229 history = box_message[:-1]230 231 if (not llm_dict):232 box_message[-1][1] = MODEL_WARNING233 return box_message, "", ""234 235 if not ref_message:236 ref_message = question237 details = f"Q: {question}"238 else:239 details = f"Q: {question}\nR: {ref_message}"240 241 242 llm = llm_dict[llm_dropdown]243 244 if DOC_1 in doc_list:245 if (not db):246 box_message[-1][1] = DOCS_WARNING247 return box_message, "", ""248 249 docs = doc_similarity(ref_message, db, top_k, score)250 delta_top_k = top_k - len(docs)251 252 if delta_top_k > 0:253 docs = doc_similarity(ref_message, db, top_k+delta_top_k, score)254 255 prompt = PROMPT_DOC256 #chain = load_qa_chain(llm, chain_type="stuff")257 258 else:259 prompt = PROMPT_BASE260 docs = []261 262 chain = LLMChain(llm = llm,263 prompt = prompt,264 output_key = 'output_text')265 266 all_output = chain({"question": question,267 "context": docs,268 "chat_history": get_chat_history(history)269 })270 271 272 bot_message = all_output['output_text']273 274 source = "".join([f"""<details> <summary>{doc.metadata["source"]}</summary>275{doc.page_content}276 277</details>""" for i, doc in enumerate(docs)])278 279 #print(source)280 281 box_message[-1][1] = bot_message282 return box_message, "", [[details, bot_message + '\n\nMetadata:\n' + source]]283 284#----------------------------------------------------------------------------------------------------------285#----------------------------------------------------------------------------------------------------------286 287with gr.Blocks(288 title = TAB_1,289 theme = "Base",290 css = """.bigbox {291 min-height:250px;292}293""") as demo:294 llm = gr.State()295 chain_2 = gr.State() # not inuse296 vector_db = gr.State()297 gr.Markdown(webui_title)298 gr.Markdown(dup_link)299 gr.Markdown(init_message)300 301 with gr.Row():302 with gr.Column(scale=10):303 llm_api_textbox = gr.Textbox(304 label = "OpenAI API Key",305 # show_label = False,306 value = OPENAI_API_KEY,307 placeholder = "Paste Your OpenAI API Key (sk-...) and Hit ENTER",308 lines=1,309 type='password')310 311 with gr.Column(scale=1, min_width=BUTTON_MIN_WIDTH):312 313 init = gr.Button(KEY_INIT) #.style(full_width=False)314 model_statusbox = gr.HTML(MODEL_NULL+DOCS_NULL)315 316 with gr.Tab(TAB_1):317 with gr.Row():318 with gr.Column(scale=10):319 chatbot = gr.Chatbot(elem_classes="bigbox")320 #with gr.Column(scale=1):321 with gr.Column(scale=1, min_width=BUTTON_MIN_WIDTH):322 doc_check = gr.CheckboxGroup(choices = DOC_SUPPORTED,323 value = DOC_DEFAULT,324 label = DOC_LABEL,325 interactive=True)326 llm_dropdown = gr.Dropdown(LLM_LIST,327 value=LLM_LIST[0],328 multiselect=False,329 interactive=True,330 label="LLM Selection",331 )332 with gr.Row():333 with gr.Column(scale=10):334 query = gr.Textbox(label="Question:",335 lines=2)336 ref = gr.Textbox(label="Reference(optional):")337 338 with gr.Column(scale=1, min_width=BUTTON_MIN_WIDTH):339 340 clear = gr.Button(KEY_CLEAR)341 submit = gr.Button(KEY_SUBMIT,variant="primary")342 343 344 with gr.Tab(TAB_2):345 with gr.Row():346 with gr.Column():347 top_k = gr.Slider(1,348 TOP_K_MAX,349 value=TOP_K_DEFAULT,350 step=1,351 label="Vector similarity top_k",352 interactive=True)353 with gr.Column():354 score = gr.Slider(0.01,355 0.99,356 value=SCORE_DEFAULT,357 step=0.01,358 label="Vector similarity score",359 interactive=True)360 detail_panel = gr.Chatbot(label="Related Docs")361 362 with gr.Tab(TAB_3):363 with gr.Row():364 with gr.Column():365 emb_textbox = gr.Textbox(366 label = "Embedding Model",367 # show_label = False,368 value = EMBEDDING_MODEL,369 placeholder = "Paste Your Embedding Model Repo on HuggingFace",370 lines=1,371 interactive=True,372 type='email')373 374 with gr.Column():375 emb_dropdown = gr.Dropdown(376 EMBEDDING_LIST,377 value=EMBEDDING_LOADER,378 multiselect=False,379 interactive=True,380 label="Embedding Loader")381 382 with gr.Accordion("Pinecone Database for "+DOC_1):383 with gr.Row():384 db_api_textbox = gr.Textbox(385 label = "Pinecone API Key",386 # show_label = False,387 value = PINECONE_KEY,388 placeholder = "Paste Your Pinecone API Key (xx-xx-xx-xx-xx) and Hit ENTER",389 lines=1,390 interactive=True,391 type='password')392 with gr.Row():393 db_env_textbox = gr.Textbox(394 label = "Pinecone Environment",395 # show_label = False,396 value = PINECONE_ENV,397 placeholder = "Paste Your Pinecone Environment (xx-xx-xx) and Hit ENTER",398 lines=1,399 interactive=True,400 type='email')401 db_index_textbox = gr.Textbox(402 label = "Pinecone Index",403 # show_label = False,404 value = PINECONE_INDEX,405 placeholder = "Paste Your Pinecone Index (xxxx) and Hit ENTER",406 lines=1,407 interactive=True,408 type='email')409 with gr.Tab(TAB_4):410 "TODO"411 412 413 414 init_input = [llm_api_textbox, emb_textbox, emb_dropdown, db_api_textbox, db_env_textbox, db_index_textbox]415 init_output = [llm_api_textbox, model_statusbox,416 llm, chain_2,417 vector_db, chatbot]418 419 llm_api_textbox.submit(init_model, init_input, init_output)420 init.click(init_model, init_input, init_output)421 422 submit.click(user,423 [query, chatbot],424 [query, chatbot],425 queue=False).then(426 bot,427 [chatbot, ref,428 llm_dropdown, llm, doc_check,429 vector_db, top_k, score],430 [chatbot, ref, detail_panel]431 )432 433 clear.click(lambda: (None,None,None), None, [query, ref, chatbot], queue=False)434 435#----------------------------------------------------------------------------------------------------------436#----------------------------------------------------------------------------------------------------------437 438if __name__ == "__main__":439 demo.launch(share = False,440 inbrowser = True,441 favicon_path = FAVICON)442 443