kyleebrooks/VectorDatabaseCreate
1
1from llama_index import SimpleDirectoryReader, Prompt, LLMPredictor, GPTVectorStoreIndex, VectorStoreIndex, PromptHelper, ServiceContext, load_index_from_storage, StorageContext2from llama_index.node_parser import SimpleNodeParser3from llama_index.data_structs import Node4from langchain_community.chat_models import ChatOpenAI5from huggingface_hub import whoami6from huggingface_hub import HfApi7from huggingface_hub import login8import os9import openai10import tiktoken11import shutil12import gradio as gr13 14 15 16#if you have OpenAI API key as a string, enable the below17openai.api_key = ""18os.environ["OPENAI_API_KEY"] = ''19large_document=""20api=HfApi()21model_type=""22messages = []23Chat_message = []24chat_history=[]25custom_chat_history=[]26max_input_size = 409627num_outputs = 51228chunk_size_limit = 60029chunk_overlap_ratio = .130 31 32prompt_helper = PromptHelper(max_input_size, num_outputs, chunk_overlap_ratio, chunk_size_limit)33 34store = './storage'35#store = 'kyleebrooks/Data/storage'36 37max_response_tokens = 100038token_limit= 409739 40template = (41 "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"42 "---------------------\n"43 "{context_str}"44 "\n---------------------\n"45 "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"46)47qa_template = Prompt(template)48 49def upload_file (index, input_file):50 HF_TOKEN = os.getenv('HF_TOKEN')51 login(token=HF_TOKEN)52 json_list=["docstore.json", "graph_store.json", "index_store.json", "vector_store.json"]53 os.mkdir("/tmp/gradio/json")54 index.storage_context.persist(persist_dir="/tmp/gradio/json") 55 for i in json_list:56 print(i)57 api.upload_file(58 path_or_fileobj="/tmp/gradio/json/"+i,59 #path_or_fileobj=i.name,60 path_in_repo="storage/"+i,61 repo_id="kyleebrooks/VectorDatabaseCreate",62 repo_type="space" # dataset63 )64 65#loads openai key66def load_api_key (api_key):67 os.environ["OPENAI_API_KEY"] = str(api_key)68 openai.api_key = str(api_key)69 70#identifies the current number of tokens used for the conversation71def num_tokens_from_messages(messages, model_type):72 encoding = tiktoken.encoding_for_model(model_type)73 num_tokens = 074 for message in messages:75 num_tokens += 4 # every message follows <im_start>{role/name}\n{content}<im_end>\n76 for key, value in message.items():77 num_tokens += len(encoding.encode(value))78 if key == "name": # if there's a name, the role is omitted79 num_tokens += -1 # role is always required and always 1 token80 num_tokens += 2 # every reply is primed with <im_start>assistant81 print(num_tokens)82 return num_tokens83 84#constructs the index and saves to a subfolder 85def construct_index(create_index, input_file, model_type, save_index):86 if create_index == "Yes":87 HF_TOKEN = os.getenv('HF_TOKEN')88 login(token=HF_TOKEN)89 source=input_file[0].name90 suffix = source.rsplit("/", 1)[1]91 prefix = source.rsplit("/", 2)[0]92 directories=[]93 print(prefix+" This is the Prefix")94 for i in input_file:95 directories.append(i.name)96 print(i.name)97 response="constructing index"98 print('Constructing index')99 # load in the documents from the docs subfolder100 llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0.1, model_name=model_type, max_tokens=num_outputs))101 service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)102 docs = SimpleDirectoryReader(input_files=directories, filename_as_id=True).load_data()103 #Large_document=str(docs)104 #node_parser = SimpleNodeParser.from_defaults(chunk_size=1024, chunk_overlap=20)105 # Use the Node Parser to get nodes from the document106 #nodes = node_parser.get_nodes_from_documents([large_document], show_progress=False)107 # Each node in the 'nodes' list will contain a smaller chunk of the text file 108 109 #index = GPTVectorStoreIndex.from_documents(nodes, service_context=service_context)110 index = GPTVectorStoreIndex.from_documents(docs, service_context=service_context)111 #index = VectorStoreIndex.from_documents(docs, service_context=service_context) 112 index.set_index_id('vector_index')113 # Stores json files in a subfolder114 if save_index=="Yes":115 upload_file(index, input_file)116 index_status="Index constructed and saved, allow time for loading"117 else:118 index_status="Index constructed but not saved for future use"119 index.storage_context.persist(persist_dir=store)120 # Clears out temporary files121 shutil.rmtree(prefix)122 response=index_status123 return response124 else:125 response= "You did not select Yes to load a new index."126 127 return response128 129 130#resets the conversation131def generate_restart(prompt, model_type):132 133 messages.clear()134 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."})135 storage_context = StorageContext.from_defaults(persist_dir=store)136 llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0.1, model_name=model_type, max_tokens=num_outputs))137 service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)138 #index = load_index_from_storage(storage_context)139 index = load_index_from_storage(140 StorageContext.from_defaults(persist_dir=store),141 service_context=service_context,142 )143 #query_engine = index.as_query_engine(text_qa_template=qa_template)144 chat_engine = index.as_chat_engine(text_qa_template=qa_template)145 string_message=str(messages)146 #response = query_engine.query(string_message)147 response = chat_engine.chat(messages)148 messages.clear()149 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."})150 messages.append({"role":"user","content": ""})151 messages.append({"role":"assistant","content": ""})152 print("restert initiated")153 print(messages)154 return response.response155 156#generates the ChatGPT call157def generate_response(prompt, model_type):158 159 messages.append({"role": "user", "content": prompt})160 storage_context = StorageContext.from_defaults(persist_dir=store)161 llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0.1, model_name=model_type, max_tokens=num_outputs))162 service_context = ServiceContext.from_defaults(llm=ChatOpenAI(temperature=0., model_name=model_type))163 #service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)164 index = load_index_from_storage(165 StorageContext.from_defaults(persist_dir=store),166 service_context=service_context,167 )168 #chat_engine = index.as_chat_engine(verbose=True, chat_history=chat_history, text_qa_template=qa_template, chat_mode='condense_question')169 query_engine = index.as_query_engine(text_qa_template=qa_template) 170 string_message=str(messages)171 response = query_engine.query(prompt)172 #response = chat_engine.chat(prompt, chat_history)173 string_response=str(response)174 messages.append({"role": "assistant", "content":string_response})175 num_tokens_from_messages(messages, model_type)176 print(messages)177 print("below is history")178 print(chat_history)179 180 return ('MIL Custom Index Chatbot: '+response.response)181 182 183def my_chatbot(input, history, model_type):184 history = history or []185 if num_tokens_from_messages(messages, model_type)<(int(token_limit)-int(max_response_tokens)):186 output = generate_response(input, model_type)187 history.append((input, output))188 return history, history189 else:190 history.clear()191 output = generate_restart(input, model_type)192 history.append((input, output))193 prompt=input194 return prompt, prompt195 196def index_chatbot(input_text):197 if not hasattr(chatbot, 'index'):198 storage_context = StorageContext.from_defaults(persist_dir=store)199 index = load_index_from_storage(storage_context)200 query_engine = chatbot.index.as_query_engine(text_qa_template=QA_TEMPLATE)201 response = chatbot.query_engine.query(input_text)202 return response.response203 204 205with gr.Blocks() as demo:206 207 gr.Markdown("""<h1><center>MIL Custom Vector Index Chatbot</center></h1>""")208 gr.Image(value="logo.PNG", width=200, height=150, interactive=False, show_share_button=False)209 api_key = gr.Textbox(type='password', label="Enter the API key", width=250)210 input_file = gr.Files()211 #load_btn.click(in_to_out,input_file,output_file)212 with gr.Row(equal_height=True):213 create_index = gr.Radio(["Yes", "No"], label = "index creation", info="Would you like to create a new index?", value="No")214 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") 215 save_index = gr.Radio(["Yes", "No"], label = "Save Index", info="Would you like to save the index for future use?", value="No")216 output = gr.Textbox(217 label="Output",218 info="",219 lines=1220 )221 submit_index = gr.Button("Create Index")222 submit_index.click(load_api_key, [api_key])223 chatbot = gr.Chatbot()224 state = gr.State()225 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.")226 submit = gr.Button("SEND")227 submit.click(load_api_key, [api_key])228 submit.click(my_chatbot, inputs=[text, state, model_type], outputs=[chatbot, state])229 submit_index.click(construct_index, [create_index, input_file, model_type, save_index], output, show_progress=True)230 231 232demo.launch(share = False)233 234 235 236