Spaces:
Running
Running
import glob | |
import gradio as gr | |
import openai | |
import os | |
from dotenv import load_dotenv | |
import phoenix as px | |
import llama_index | |
from llama_index import Prompt, ServiceContext, VectorStoreIndex, SimpleDirectoryReader | |
from llama_index.chat_engine.types import ChatMode | |
from llama_index.llms import ChatMessage, MessageRole | |
from llama_index.vector_stores.qdrant import QdrantVectorStore | |
from llama_index.text_splitter import SentenceSplitter | |
from llama_index.extractors import TitleExtractor | |
from llama_index.ingestion import IngestionPipeline | |
from chat_template import CHAT_TEXT_QA_PROMPT, TEXT_QA_SYSTEM_PROMPT | |
from schemas import ChatbotVersion, ServiceProvider | |
from chatbot import Chatbot, IndexBuilder | |
from custom_io import MarkdownReader, UnstructuredReader, default_file_metadata_func | |
from qdrant import client as qdrantClient | |
from llama_index import set_global_service_context | |
from service_provider_config import get_service_provider_config | |
load_dotenv() | |
# initial service setup | |
px.launch_app() | |
llama_index.set_global_handler("arize_phoenix") | |
# llama_index.set_global_handler("wandb", run_args={"project": "llamaindex"}) | |
openai.api_key = os.getenv("OPENAI_API_KEY") | |
IS_LOAD_FROM_VECTOR_STORE = True | |
VDB_COLLECTION_NAME = "demo-v5" | |
MODEL_NAME = ChatbotVersion.CHATGPT_4.value | |
CHUNK_SIZE = 8191 | |
LLM, EMBED_MODEL = get_service_provider_config( | |
service_provider=ServiceProvider.OPENAI, model_name=MODEL_NAME) | |
service_context = ServiceContext.from_defaults( | |
chunk_size=CHUNK_SIZE, | |
llm=LLM, | |
embed_model=EMBED_MODEL, | |
) | |
set_global_service_context(service_context) | |
class AwesumIndexBuilder(IndexBuilder): | |
def _load_doucments(self): | |
directory = "./awesumcare_data/awesumcare_manual_data" | |
dir_reader = SimpleDirectoryReader(directory, file_extractor={ | |
".pdf": UnstructuredReader(), | |
".docx": UnstructuredReader(), | |
".pptx": UnstructuredReader(), | |
".md": MarkdownReader() | |
}, | |
recursive=True, | |
exclude=["*.png", "*.pptx", "*.docx", "*.pdf"], | |
file_metadata=default_file_metadata_func) | |
self.documents = dir_reader.load_data() | |
print(f"Loaded {len(self.documents)} docs") | |
def _setup_service_context(self): | |
super()._setup_service_context() | |
def _setup_vector_store(self): | |
self.vector_store = QdrantVectorStore( | |
client=qdrantClient, collection_name=self.vdb_collection_name) | |
super()._setup_vector_store() | |
def _setup_index(self): | |
super()._setup_index() | |
if self.is_load_from_vector_store: | |
self.index = VectorStoreIndex.from_vector_store(self.vector_store) | |
print("set up index from vector store") | |
return | |
pipeline = IngestionPipeline( | |
transformations=[ | |
self.embed_model, | |
], | |
vector_store=self.vector_store, | |
) | |
pipeline.run(documents=self.documents, show_progress=True) | |
self.index = VectorStoreIndex.from_vector_store(self.vector_store) | |
class AwesumCareToolChatbot(Chatbot): | |
DENIED_ANSWER_PROMPT = "" | |
SYSTEM_PROMPT = "" | |
CHAT_EXAMPLES = [ | |
"什麼是安心三寶?", | |
"點樣立平安紙?", | |
"甚麼是⾒證?", | |
"訂立每份⽂件需要多少錢以及付款⽅法?", | |
"通過安⼼三寶製作的⽂件有法律效⼒嗎?", | |
] | |
def _setup_observer(self): | |
pass | |
def _setup_index(self): | |
super()._setup_index() | |
def _setup_query_engine(self): | |
super()._setup_query_engine() | |
self.query_engine = self.index.as_query_engine( | |
text_qa_template=CHAT_TEXT_QA_PROMPT) | |
def _setup_tools(self): | |
from llama_index.tools.query_engine import QueryEngineTool | |
self.tools = QueryEngineTool.from_defaults( | |
query_engine=self.query_engine) | |
return super()._setup_tools() | |
def _setup_chat_engine(self): | |
from llama_index.agent import OpenAIAgent | |
self.chat_engine = OpenAIAgent.from_tools( | |
tools=[self.tools], | |
llm=LLM, | |
similarity_top_k=1, | |
verbose=True | |
) | |
print("set up agent as chat engine") | |
super()._setup_chat_engine() | |
class AweSumCareContextChatbot(AwesumCareToolChatbot): | |
def _setup_query_engine(self): | |
pass | |
def _setup_tools(self): | |
pass | |
def _setup_chat_engine(self): | |
self.chat_engine = self.index.as_chat_engine( | |
chat_mode=ChatMode.CONTEXT, | |
similarity_top_k=5, | |
system_prompt=TEXT_QA_SYSTEM_PROMPT.content, | |
text_qa_template=CHAT_TEXT_QA_PROMPT) | |
class AweSumCareSimpleChatbot(AwesumCareToolChatbot): | |
def _setup_query_engine(self): | |
pass | |
def _setup_tools(self): | |
pass | |
def _setup_chat_engine(self): | |
self.chat_engine = self.index.as_chat_engine( | |
chat_mode=ChatMode.SIMPLE) | |
model_name = MODEL_NAME | |
index_builder = AwesumIndexBuilder(vdb_collection_name=VDB_COLLECTION_NAME, | |
embed_model=EMBED_MODEL, | |
is_load_from_vector_store=IS_LOAD_FROM_VECTOR_STORE) | |
# gpt-3.5-turbo-1106, gpt-4-1106-preview | |
awesum_chatbot = AwesumCareToolChatbot(model_name=model_name, index_builder=index_builder) | |
awesum_chatbot_context = AweSumCareContextChatbot(model_name=model_name, index_builder=index_builder) | |
awesum_chatbot_simple = AweSumCareSimpleChatbot(model_name=model_name, index_builder=index_builder) | |
def service_setup(model_name): | |
CHUNK_SIZE = 1024 | |
LLM, EMBED_MODEL = get_service_provider_config( | |
service_provider=ServiceProvider.OPENAI, model_name=model_name) | |
service_context = ServiceContext.from_defaults( | |
chunk_size=CHUNK_SIZE, | |
llm=LLM, | |
embed_model=EMBED_MODEL, | |
) | |
set_global_service_context(service_context) | |
return LLM, EMBED_MODEL | |
chatbot = gr.Chatbot(height=500) | |
def vote(data: gr.LikeData): | |
if data.liked: | |
print("You upvoted this response: " + data.value) | |
else: | |
print("You downvoted this response: " + data.value) | |
with gr.Blocks(fill_height=True) as demo: | |
gr.Markdown("# Awesum Care demo") | |
with gr.Tab("With relevant context sent to system prompt"): | |
context_interface = gr.ChatInterface( | |
awesum_chatbot_context.stream_chat, | |
examples=awesum_chatbot.CHAT_EXAMPLES, | |
chatbot=chatbot | |
) | |
# gr.Markdown("instructions:\n" | |
# "\nUsing model gpt-4-preview-1106, the most advanced model now in the market.\n" | |
# "\n(Note that it can be much slower than gpt-3.5, openai's api can be unstable sometimes.)\n" | |
# ) | |
demo.queue() | |
demo.launch(share=False, auth=("demo", os.getenv("PASSWORD"))) | |