from llama_index import SimpleDirectoryReader, Prompt, LLMPredictor, GPTVectorStoreIndex, VectorStoreIndex, PromptHelper, ServiceContext, load_index_from_storage, StorageContext from llama_index.node_parser import SimpleNodeParser from llama_index.data_structs import Node from langchain_community.chat_models import ChatOpenAI from huggingface_hub import whoami from huggingface_hub import HfApi from huggingface_hub import login import os import openai import tiktoken import shutil import gradio as gr #if you have OpenAI API key as a string, enable the below openai.api_key = "" os.environ["OPENAI_API_KEY"] = '' large_document="" api=HfApi() model_type="" messages = [] Chat_message = [] chat_history=[] custom_chat_history=[] max_input_size = 4096 num_outputs = 512 chunk_size_limit = 600 chunk_overlap_ratio = .1 prompt_helper = PromptHelper(max_input_size, num_outputs, chunk_overlap_ratio, chunk_size_limit) store = './storage' #store = 'kyleebrooks/Data/storage' max_response_tokens = 1000 token_limit= 4097 template = ( "This Chatbot is helpful, accurate, and will use the context below for answering all questions. This Chatbot will not answer questions not included in the context provided \n" "---------------------\n" "{context_str}" "\n---------------------\n" "Given this information, please answer the question by providing a detailed summary and provide accurate citations for all referenced areas at the end of each response. {query_str}\n" ) qa_template = Prompt(template) def upload_file (index, input_file): HF_TOKEN = os.getenv('HF_TOKEN') login(token=HF_TOKEN) json_list=["docstore.json", "graph_store.json", "index_store.json", "vector_store.json"] os.mkdir("/tmp/gradio/json") index.storage_context.persist(persist_dir="/tmp/gradio/json") for i in json_list: print(i) api.upload_file( path_or_fileobj="/tmp/gradio/json/"+i, #path_or_fileobj=i.name, path_in_repo="storage/"+i, repo_id="kyleebrooks/VectorDatabaseCreate", repo_type="space" # dataset ) #loads openai key def load_api_key (api_key): os.environ["OPENAI_API_KEY"] = str(api_key) openai.api_key = str(api_key) #identifies the current number of tokens used for the conversation def num_tokens_from_messages(messages, model_type): encoding = tiktoken.encoding_for_model(model_type) num_tokens = 0 for message in messages: num_tokens += 4 # every message follows {role/name}\n{content}\n for key, value in message.items(): num_tokens += len(encoding.encode(value)) if key == "name": # if there's a name, the role is omitted num_tokens += -1 # role is always required and always 1 token num_tokens += 2 # every reply is primed with assistant print(num_tokens) return num_tokens #constructs the index and saves to a subfolder def construct_index(create_index, input_file, model_type, save_index): if create_index == "Yes": HF_TOKEN = os.getenv('HF_TOKEN') login(token=HF_TOKEN) source=input_file[0].name suffix = source.rsplit("/", 1)[1] prefix = source.rsplit("/", 2)[0] directories=[] print(prefix+" This is the Prefix") for i in input_file: directories.append(i.name) print(i.name) response="constructing index" print('Constructing index') # load in the documents from the docs subfolder llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0.1, model_name=model_type, max_tokens=num_outputs)) service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper) docs = SimpleDirectoryReader(input_files=directories, filename_as_id=True).load_data() #Large_document=str(docs) #node_parser = SimpleNodeParser.from_defaults(chunk_size=1024, chunk_overlap=20) # Use the Node Parser to get nodes from the document #nodes = node_parser.get_nodes_from_documents([large_document], show_progress=False) # Each node in the 'nodes' list will contain a smaller chunk of the text file #index = GPTVectorStoreIndex.from_documents(nodes, service_context=service_context) index = GPTVectorStoreIndex.from_documents(docs, service_context=service_context) #index = VectorStoreIndex.from_documents(docs, service_context=service_context) index.set_index_id('vector_index') # Stores json files in a subfolder if save_index=="Yes": upload_file(index, input_file) index_status="Index constructed and saved, allow time for loading" else: index_status="Index constructed but not saved for future use" index.storage_context.persist(persist_dir=store) # Clears out temporary files shutil.rmtree(prefix) response=index_status return response else: response= "You did not select Yes to load a new index." return response #resets the conversation def generate_restart(prompt, model_type): messages.clear() messages.append({"role":"system", "content": "Tell the user that this conversation has been reset due to the discussion size reaching maximum size, and to please start by asking a new question."}) storage_context = StorageContext.from_defaults(persist_dir=store) llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0.1, model_name=model_type, max_tokens=num_outputs)) service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper) #index = load_index_from_storage(storage_context) index = load_index_from_storage( StorageContext.from_defaults(persist_dir=store), service_context=service_context, ) #query_engine = index.as_query_engine(text_qa_template=qa_template) chat_engine = index.as_chat_engine(text_qa_template=qa_template) string_message=str(messages) #response = query_engine.query(string_message) response = chat_engine.chat(messages) messages.clear() messages.append({"role":"system", "content": "This Chatbot is helpful, accurate, and provides all relevnt information from the Treasury Financial Manual (TFM) when responding. This Chatbot always provides accurate citations from the TFM."}) messages.append({"role":"user","content": ""}) messages.append({"role":"assistant","content": ""}) print("restert initiated") print(messages) return response.response #generates the ChatGPT call def generate_response(prompt, model_type): messages.append({"role": "user", "content": prompt}) storage_context = StorageContext.from_defaults(persist_dir=store) llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0.1, model_name=model_type, max_tokens=num_outputs)) service_context = ServiceContext.from_defaults(llm=ChatOpenAI(temperature=0., model_name=model_type)) #service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper) index = load_index_from_storage( StorageContext.from_defaults(persist_dir=store), service_context=service_context, ) #chat_engine = index.as_chat_engine(verbose=True, chat_history=chat_history, text_qa_template=qa_template, chat_mode='condense_question') query_engine = index.as_query_engine(text_qa_template=qa_template) string_message=str(messages) response = query_engine.query(prompt) #response = chat_engine.chat(prompt, chat_history) string_response=str(response) messages.append({"role": "assistant", "content":string_response}) num_tokens_from_messages(messages, model_type) print(messages) print("below is history") print(chat_history) return ('MIL Custom Index Chatbot: '+response.response) def my_chatbot(input, history, model_type): history = history or [] if num_tokens_from_messages(messages, model_type)<(int(token_limit)-int(max_response_tokens)): output = generate_response(input, model_type) history.append((input, output)) return history, history else: history.clear() output = generate_restart(input, model_type) history.append((input, output)) prompt=input return prompt, prompt def index_chatbot(input_text): if not hasattr(chatbot, 'index'): storage_context = StorageContext.from_defaults(persist_dir=store) index = load_index_from_storage(storage_context) query_engine = chatbot.index.as_query_engine(text_qa_template=QA_TEMPLATE) response = chatbot.query_engine.query(input_text) return response.response with gr.Blocks() as demo: gr.Markdown("""

MIL Custom Vector Index Chatbot

""") gr.Image(value="logo.PNG", width=200, height=150, interactive=False, show_share_button=False) api_key = gr.Textbox(type='password', label="Enter the API key", width=250) input_file = gr.Files() #load_btn.click(in_to_out,input_file,output_file) with gr.Row(equal_height=True): create_index = gr.Radio(["Yes", "No"], label = "index creation", info="Would you like to create a new index?", value="No") model_type = gr.Radio(["gpt-3.5-turbo", "gpt-4"], label = "Model_Type", info="Would you like to create a new index?", value="gpt-3.5-turbo") save_index = gr.Radio(["Yes", "No"], label = "Save Index", info="Would you like to save the index for future use?", value="No") output = gr.Textbox( label="Output", info="", lines=1 ) submit_index = gr.Button("Create Index") submit_index.click(load_api_key, [api_key]) chatbot = gr.Chatbot() state = gr.State() text = gr.Textbox(label="Input", info="", lines=2, placeholder="Hello. Ask me a question about the indexed content. Please approach each question as if it is a new question, my memory is limited in this model.") submit = gr.Button("SEND") submit.click(load_api_key, [api_key]) submit.click(my_chatbot, inputs=[text, state, model_type], outputs=[chatbot, state]) submit_index.click(construct_index, [create_index, input_file, model_type, save_index], output, show_progress=True) demo.launch(share = False)