Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import random | |
| import time | |
| from langchain import PromptTemplate | |
| from langchain.llms import OpenAI | |
| from langchain.chat_models import ChatOpenAI | |
| from langchain.embeddings import HuggingFaceEmbeddings, HuggingFaceInstructEmbeddings, OpenAIEmbeddings | |
| from langchain.vectorstores import Pinecone | |
| from langchain.chains import LLMChain | |
| from langchain.chains.question_answering import load_qa_chain | |
| import pinecone | |
| import os | |
| os.environ["TOKENIZERS_PARALLELISM"] = "false" | |
| #OPENAI_API_KEY = "" | |
| OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "") | |
| OPENAI_TEMP = 1 | |
| OPENAI_API_LINK = "[OpenAI API Key](https://platform.openai.com/account/api-keys)" | |
| OPENAI_LINK = "[OpenAI](https://openai.com)" | |
| PINECONE_KEY = os.environ.get("PINECONE_KEY", "") | |
| PINECONE_ENV = os.environ.get("PINECONE_ENV", "asia-northeast1-gcp") | |
| PINECONE_INDEX = os.environ.get("PINECONE_INDEX", '3gpp-r16') | |
| PINECONE_LINK = "[Pinecone](https://www.pinecone.io)" | |
| LANGCHAIN_LINK = "[LangChain](https://python.langchain.com/en/latest/index.html)" | |
| EMBEDDING_MODEL = os.environ.get("EMBEDDING_MODEL", "hkunlp/instructor-large") | |
| EMBEDDING_LOADER = os.environ.get("EMBEDDING_LOADER", "HuggingFaceInstructEmbeddings") | |
| EMBEDDING_LIST = ["HuggingFaceInstructEmbeddings", "HuggingFaceEmbeddings", "OpenAIEmbeddings"] | |
| # return top-k text chunks from vector store | |
| TOP_K_DEFAULT = 15 | |
| TOP_K_MAX = 30 | |
| SCORE_DEFAULT = 0.33 | |
| BUTTON_MIN_WIDTH = 215 | |
| LLM_NULL = "LLM-UNLOAD-critical" | |
| LLM_DONE = "LLM-LOADED-9cf" | |
| DB_NULL = "DB-UNLOAD-critical" | |
| DB_DONE = "DB-LOADED-9cf" | |
| FORK_BADGE = "Fork-HuggingFace Space-9cf" | |
| def get_logo(inputs, logo) -> str: | |
| return f"""https://img.shields.io/badge/{inputs}?style=flat&logo={logo}&logoColor=white""" | |
| def get_status(inputs, logo, pos) -> str: | |
| return f"""<img | |
| src = "{get_logo(inputs, logo)}"; | |
| style = "margin: 0 auto;float:{pos};border: 2px solid transparent;"; | |
| >""" | |
| KEY_INIT = "Initialize Model" | |
| KEY_SUBMIT = "Submit" | |
| KEY_CLEAR = "Clear" | |
| MODEL_NULL = get_status(LLM_NULL, "openai", "right") | |
| MODEL_DONE = get_status(LLM_DONE, "openai", "right") | |
| DOCS_NULL = get_status(DB_NULL, "processingfoundation", "right") | |
| DOCS_DONE = get_status(DB_DONE, "processingfoundation", "right") | |
| TAB_1 = "Chatbot" | |
| TAB_2 = "Details" | |
| TAB_3 = "Database" | |
| TAB_4 = "TODO" | |
| FAVICON = './icon.svg' | |
| LLM_LIST = ["gpt-3.5-turbo", "text-davinci-003"] | |
| DOC_1 = '3GPP' | |
| DOC_2 = 'HTTP2' | |
| DOC_SUPPORTED = [DOC_1] | |
| DOC_DEFAULT = [DOC_1] | |
| DOC_LABEL = "Reference Docs" | |
| MODEL_WARNING = f"Please paste your **{OPENAI_API_LINK}** and then **{KEY_INIT}**" | |
| DOCS_WARNING = f"""Database Unloaded | |
| Please check your **{TAB_3}** config and then **{KEY_INIT}** | |
| Or you could uncheck **{DOC_LABEL}** to ask LLM directly""" | |
| webui_title = """ | |
| # OpenAI Chatbot Based on Vector Database | |
| """ | |
| dup_link = f'''<a href="https://huggingface.co/spaces/ShawnAI/VectorDB-ChatBot?duplicate=true" | |
| style="display:grid; width: 200px;"> | |
| <img src="{get_logo(FORK_BADGE, "addthis")}"></a>''' | |
| init_message = f"""This demonstration website is based on \ | |
| **{OPENAI_LINK}** with **{LANGCHAIN_LINK}** and **{PINECONE_LINK}** | |
| 1. Insert your **{OPENAI_API_LINK}** and click `{KEY_INIT}` | |
| 2. Insert your **Question** and click `{KEY_SUBMIT}` | |
| """ | |
| PROMPT_DOC = PromptTemplate( | |
| input_variables=["context", "chat_history", "question"], | |
| template="""Context: | |
| ## | |
| {context} | |
| ## | |
| Chat History: | |
| ## | |
| {chat_history} | |
| ## | |
| Question: | |
| {question} | |
| Answer:""" | |
| ) | |
| PROMPT_BASE = PromptTemplate( | |
| input_variables=['question', "chat_history"], | |
| template="""Chat History: | |
| ## | |
| {chat_history} | |
| ## | |
| Question: | |
| ## | |
| {question} | |
| ## | |
| Answer:""" | |
| ) | |
| #---------------------------------------------------------------------------------------------------------- | |
| #---------------------------------------------------------------------------------------------------------- | |
| def init_rwkv(): | |
| try: | |
| import rwkv | |
| return True | |
| except Exception: | |
| print("RWKV not found, skip local llm") | |
| return False | |
| def init_model(api_key, emb_name, emb_loader, db_api_key, db_env, db_index): | |
| init_rwkv() | |
| try: | |
| if not (api_key and api_key.startswith("sk-") and len(api_key) > 50): | |
| return None,MODEL_NULL+DOCS_NULL,None,None,None,None | |
| llm_dict = {} | |
| for llm_name in LLM_LIST: | |
| if llm_name == "gpt-3.5-turbo": | |
| llm_dict[llm_name] = ChatOpenAI(model_name=llm_name, | |
| temperature = OPENAI_TEMP, | |
| openai_api_key = api_key | |
| ) | |
| else: | |
| llm_dict[llm_name] = OpenAI(model_name=llm_name, | |
| temperature = OPENAI_TEMP, | |
| openai_api_key = api_key) | |
| if not (emb_name and db_api_key and db_env and db_index): | |
| return api_key,MODEL_DONE+DOCS_NULL,llm_dict,None,None,None | |
| if emb_loader == "OpenAIEmbeddings": | |
| embeddings = eval(emb_loader)(openai_api_key=api_key) | |
| else: | |
| embeddings = eval(emb_loader)(model_name=emb_name) | |
| pinecone.init(api_key = db_api_key, | |
| environment = db_env) | |
| db = Pinecone.from_existing_index(index_name = db_index, | |
| embedding = embeddings) | |
| return api_key, MODEL_DONE+DOCS_DONE, llm_dict, None, db, None | |
| except Exception as e: | |
| print(e) | |
| return None,MODEL_NULL+DOCS_NULL,None,None,None,None | |
| def get_chat_history(inputs) -> str: | |
| res = [] | |
| for human, ai in inputs: | |
| res.append(f"Q: {human}\nA: {ai}") | |
| return "\n".join(res) | |
| def remove_duplicates(documents, score_min): | |
| seen_content = set() | |
| unique_documents = [] | |
| for (doc, score) in documents: | |
| if (doc.page_content not in seen_content) and (score >= score_min): | |
| seen_content.add(doc.page_content) | |
| unique_documents.append(doc) | |
| return unique_documents | |
| def doc_similarity(query, db, top_k, score): | |
| docs = db.similarity_search_with_score(query = query, | |
| k=top_k) | |
| #docsearch = db.as_retriever(search_kwargs={'k':top_k}) | |
| #docs = docsearch.get_relevant_documents(query) | |
| udocs = remove_duplicates(docs, score) | |
| return udocs | |
| def user(user_message, history): | |
| return "", history+[[user_message, None]] | |
| def bot(box_message, ref_message, | |
| llm_dropdown, llm_dict, doc_list, | |
| db, top_k, score): | |
| # bot_message = random.choice(["Yes", "No"]) | |
| # 0 is user question, 1 is bot response | |
| question = box_message[-1][0] | |
| history = box_message[:-1] | |
| if (not llm_dict): | |
| box_message[-1][1] = MODEL_WARNING | |
| return box_message, "", "" | |
| if not ref_message: | |
| ref_message = question | |
| details = f"Q: {question}" | |
| else: | |
| details = f"Q: {question}\nR: {ref_message}" | |
| llm = llm_dict[llm_dropdown] | |
| if DOC_1 in doc_list: | |
| if (not db): | |
| box_message[-1][1] = DOCS_WARNING | |
| return box_message, "", "" | |
| docs = doc_similarity(ref_message, db, top_k, score) | |
| delta_top_k = top_k - len(docs) | |
| if delta_top_k > 0: | |
| docs = doc_similarity(ref_message, db, top_k+delta_top_k, score) | |
| prompt = PROMPT_DOC | |
| #chain = load_qa_chain(llm, chain_type="stuff") | |
| else: | |
| prompt = PROMPT_BASE | |
| docs = [] | |
| chain = LLMChain(llm = llm, | |
| prompt = prompt, | |
| output_key = 'output_text') | |
| all_output = chain({"question": question, | |
| "context": docs, | |
| "chat_history": get_chat_history(history) | |
| }) | |
| bot_message = all_output['output_text'] | |
| source = "".join([f"""<details> <summary>{doc.metadata["source"]}</summary> | |
| {doc.page_content} | |
| </details>""" for i, doc in enumerate(docs)]) | |
| #print(source) | |
| box_message[-1][1] = bot_message | |
| return box_message, "", [[details, bot_message + '\n\nMetadata:\n' + source]] | |
| #---------------------------------------------------------------------------------------------------------- | |
| #---------------------------------------------------------------------------------------------------------- | |
| with gr.Blocks( | |
| title = TAB_1, | |
| theme = "Base", | |
| css = """.bigbox { | |
| min-height:250px; | |
| } | |
| """) as demo: | |
| llm = gr.State() | |
| chain_2 = gr.State() # not inuse | |
| vector_db = gr.State() | |
| gr.Markdown(webui_title) | |
| gr.Markdown(dup_link) | |
| gr.Markdown(init_message) | |
| with gr.Row(): | |
| with gr.Column(scale=10): | |
| llm_api_textbox = gr.Textbox( | |
| label = "OpenAI API Key", | |
| # show_label = False, | |
| value = OPENAI_API_KEY, | |
| placeholder = "Paste Your OpenAI API Key (sk-...) and Hit ENTER", | |
| lines=1, | |
| type='password') | |
| with gr.Column(scale=1, min_width=BUTTON_MIN_WIDTH): | |
| init = gr.Button(KEY_INIT) #.style(full_width=False) | |
| model_statusbox = gr.HTML(MODEL_NULL+DOCS_NULL) | |
| with gr.Tab(TAB_1): | |
| with gr.Row(): | |
| with gr.Column(scale=10): | |
| chatbot = gr.Chatbot(elem_classes="bigbox") | |
| #with gr.Column(scale=1): | |
| with gr.Column(scale=1, min_width=BUTTON_MIN_WIDTH): | |
| doc_check = gr.CheckboxGroup(choices = DOC_SUPPORTED, | |
| value = DOC_DEFAULT, | |
| label = DOC_LABEL, | |
| interactive=True) | |
| llm_dropdown = gr.Dropdown(LLM_LIST, | |
| value=LLM_LIST[0], | |
| multiselect=False, | |
| interactive=True, | |
| label="LLM Selection", | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=10): | |
| query = gr.Textbox(label="Question:", | |
| lines=2) | |
| ref = gr.Textbox(label="Reference(optional):") | |
| with gr.Column(scale=1, min_width=BUTTON_MIN_WIDTH): | |
| clear = gr.Button(KEY_CLEAR) | |
| submit = gr.Button(KEY_SUBMIT,variant="primary") | |
| with gr.Tab(TAB_2): | |
| with gr.Row(): | |
| with gr.Column(): | |
| top_k = gr.Slider(1, | |
| TOP_K_MAX, | |
| value=TOP_K_DEFAULT, | |
| step=1, | |
| label="Vector similarity top_k", | |
| interactive=True) | |
| with gr.Column(): | |
| score = gr.Slider(0.01, | |
| 0.99, | |
| value=SCORE_DEFAULT, | |
| step=0.01, | |
| label="Vector similarity score", | |
| interactive=True) | |
| detail_panel = gr.Chatbot(label="Related Docs") | |
| with gr.Tab(TAB_3): | |
| with gr.Row(): | |
| with gr.Column(): | |
| emb_textbox = gr.Textbox( | |
| label = "Embedding Model", | |
| # show_label = False, | |
| value = EMBEDDING_MODEL, | |
| placeholder = "Paste Your Embedding Model Repo on HuggingFace", | |
| lines=1, | |
| interactive=True, | |
| type='email') | |
| with gr.Column(): | |
| emb_dropdown = gr.Dropdown( | |
| EMBEDDING_LIST, | |
| value=EMBEDDING_LOADER, | |
| multiselect=False, | |
| interactive=True, | |
| label="Embedding Loader") | |
| with gr.Accordion("Pinecone Database for "+DOC_1): | |
| with gr.Row(): | |
| db_api_textbox = gr.Textbox( | |
| label = "Pinecone API Key", | |
| # show_label = False, | |
| value = PINECONE_KEY, | |
| placeholder = "Paste Your Pinecone API Key (xx-xx-xx-xx-xx) and Hit ENTER", | |
| lines=1, | |
| interactive=True, | |
| type='password') | |
| with gr.Row(): | |
| db_env_textbox = gr.Textbox( | |
| label = "Pinecone Environment", | |
| # show_label = False, | |
| value = PINECONE_ENV, | |
| placeholder = "Paste Your Pinecone Environment (xx-xx-xx) and Hit ENTER", | |
| lines=1, | |
| interactive=True, | |
| type='email') | |
| db_index_textbox = gr.Textbox( | |
| label = "Pinecone Index", | |
| # show_label = False, | |
| value = PINECONE_INDEX, | |
| placeholder = "Paste Your Pinecone Index (xxxx) and Hit ENTER", | |
| lines=1, | |
| interactive=True, | |
| type='email') | |
| with gr.Tab(TAB_4): | |
| "TODO" | |
| init_input = [llm_api_textbox, emb_textbox, emb_dropdown, db_api_textbox, db_env_textbox, db_index_textbox] | |
| init_output = [llm_api_textbox, model_statusbox, | |
| llm, chain_2, | |
| vector_db, chatbot] | |
| llm_api_textbox.submit(init_model, init_input, init_output) | |
| init.click(init_model, init_input, init_output) | |
| submit.click(user, | |
| [query, chatbot], | |
| [query, chatbot], | |
| queue=False).then( | |
| bot, | |
| [chatbot, ref, | |
| llm_dropdown, llm, doc_check, | |
| vector_db, top_k, score], | |
| [chatbot, ref, detail_panel] | |
| ) | |
| clear.click(lambda: (None,None,None), None, [query, ref, chatbot], queue=False) | |
| #---------------------------------------------------------------------------------------------------------- | |
| #---------------------------------------------------------------------------------------------------------- | |
| if __name__ == "__main__": | |
| demo.launch(share = False, | |
| inbrowser = True, | |
| favicon_path = FAVICON) | |