| import logging | |
| import os | |
| import json | |
| from llama_index.indices.document_summary import DocumentSummaryIndexEmbeddingRetriever | |
| from llama_index.indices.vector_store import VectorIndexRetriever | |
| from llama_index.llms import OpenAI | |
| from llama_index.query_engine import RetrieverQueryEngine | |
| logging.basicConfig(format='%(asctime)s %(levelname)s:%(message)s', level=os.environ.get("LOGLEVEL", "INFO")) | |
| import gradio as gr | |
| from llama_index import VectorStoreIndex, StorageContext, download_loader, load_index_from_storage, ServiceContext, \ | |
| get_response_synthesizer | |
| cache = {} | |
| chatgpt = OpenAI(temperature=0, model="gpt-3.5-turbo") | |
| service_context = ServiceContext.from_defaults(llm=chatgpt, chunk_size=1024) | |
| def load_mapping_from_json(filepath): | |
| if not os.path.exists(filepath): | |
| return {} | |
| with open(filepath, 'r') as file: | |
| return json.load(file) | |
| userIdMapping = load_mapping_from_json('user_id_to_name_mapping.json') | |
| def loadData(): | |
| index_root = "./index" | |
| directory_names = os.listdir(index_root) | |
| for directory in directory_names: | |
| if os.path.isdir(f"{index_root}/{directory}"): | |
| print("Loading from existing index " + directory) | |
| storage_context = StorageContext.from_defaults(persist_dir=f"{index_root}/{directory}") | |
| index = load_index_from_storage(storage_context) | |
| vector_retriever = VectorIndexRetriever(index=index, similarity_top_k=5) | |
| response_synthesizer = get_response_synthesizer(service_context=service_context, | |
| response_mode="tree_summarize") | |
| query_engine = RetrieverQueryEngine( | |
| retriever=vector_retriever, | |
| response_synthesizer=response_synthesizer, | |
| ) | |
| cache[directory] = query_engine | |
| def chatbot(indexName, input_text): | |
| """ | |
| Chatbot function that takes in a prompt and returns a response | |
| """ | |
| query = "This data contains updates from multiple teams and sprints - "+ input_text + " Ignore headers in the content." | |
| response = cache[indexName].query(query) | |
| answer = response.response.replace('Based on the given context information', 'Based on the available information') | |
| for userId in userIdMapping: | |
| answer = answer.replace(userId, userIdMapping[userId]) | |
| return answer | |
| def main(): | |
| loadData() | |
| iface = gr.Interface(fn=chatbot, inputs=[ | |
| gr.Dropdown(cache.keys(), | |
| type="value", value="sos", label="Select Channel"), | |
| gr.Textbox(lines=7, label="Ask any question", placeholder='What are the key topics?')], outputs=gr.Textbox(lines=11, label="Response"), | |
| title="NLP Demo for Slack Data") | |
| if 'LOGIN_PASS' in os.environ: | |
| iface.launch(auth=('axiamatic', os.environ['LOGIN_PASS']), | |
| auth_message='For access, please check my Slack profile or contact me in Slack.', | |
| share=False) | |
| else: | |
| iface.launch(share=False) | |
| main() | |