File size: 3,026 Bytes
dddb40f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
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()
|