nxphi47/MultiPurpose-Chatbot-DEMO
1
1import os2from gradio.themes import ThemeClass as Theme3import numpy as np4import argparse5import gradio as gr6from typing import Any, Iterator7from typing import Iterator, List, Optional, Tuple8import filelock9import glob10import json11import time12from gradio.routes import Request13from gradio.utils import SyncToAsyncIterator, async_iteration14from gradio.helpers import special_args15import anyio16from typing import AsyncGenerator, Callable, Literal, Union, cast, Generator17 18from gradio_client.documentation import document, set_documentation_group19from gradio.components import Button, Component20from gradio.events import Dependency, EventListenerMethod21from typing import List, Optional, Union, Dict, Tuple22from tqdm.auto import tqdm23from huggingface_hub import snapshot_download24from gradio.themes import ThemeClass as Theme25 26from .base_demo import register_demo, get_demo_class, BaseDemo27 28import inspect29from typing import AsyncGenerator, Callable, Literal, Union, cast30 31import anyio32from gradio_client import utils as client_utils33from gradio_client.documentation import document34 35from gradio.blocks import Blocks36from gradio.components import (37 Button,38 Chatbot,39 Component,40 Markdown,41 State,42 Textbox,43 get_component_instance,44)45from gradio.events import Dependency, on46from gradio.helpers import create_examples as Examples # noqa: N81247from gradio.helpers import special_args48from gradio.layouts import Accordion, Group, Row49from gradio.routes import Request50from gradio.themes import ThemeClass as Theme51from gradio.utils import SyncToAsyncIterator, async_iteration52 53 54from ..globals import MODEL_ENGINE, RAG_CURRENT_FILE, RAG_EMBED, load_embeddings, get_rag_embeddings55 56from .chat_interface import (57 SYSTEM_PROMPT,58 MODEL_NAME,59 MAX_TOKENS,60 TEMPERATURE,61 CHAT_EXAMPLES,62 gradio_history_to_openai_conversations,63 gradio_history_to_conversation_prompt,64 DATETIME_FORMAT,65 get_datetime_string,66 format_conversation,67 chat_response_stream_multiturn_engine,68 ChatInterfaceDemo,69 CustomizedChatInterface,70)71 72from ..configs import (73 CHUNK_SIZE,74 CHUNK_OVERLAP,75 RAG_EMBED_MODEL_NAME,76)77 78RAG_CURRENT_VECTORSTORE = None79 80 81def load_document_split_vectorstore(file_path):82 global RAG_CURRENT_FILE, RAG_EMBED, RAG_CURRENT_VECTORSTORE83 from langchain.text_splitter import RecursiveCharacterTextSplitter84 from langchain_community.embeddings import HuggingFaceEmbeddings, HuggingFaceBgeEmbeddings85 from langchain_community.vectorstores import Chroma, FAISS86 from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader, TextLoader87 splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)88 if file_path.endswith('.pdf'):89 loader = PyPDFLoader(file_path)90 elif file_path.endswith('.docx'):91 loader = Docx2txtLoader(file_path)92 elif file_path.endswith('.txt'):93 loader = TextLoader(file_path)94 splits = loader.load_and_split(splitter)95 RAG_CURRENT_VECTORSTORE = FAISS.from_texts(texts=[s.page_content for s in splits], embedding=get_rag_embeddings())96 return RAG_CURRENT_VECTORSTORE97 98def docs_to_context_content(docs: List[Any]):99 content = "\n".join([d.page_content for d in docs])100 return content101 102 103DOC_TEMPLATE = """###104{content}105###106 107"""108 109DOC_INSTRUCTION = """Answer the following query exclusively based on the information provided in the document above. \110If the information is not found, please say so instead of making up facts! Remember to answer the question in the same language as the user query!111"""112 113 114def docs_to_rag_context(docs: List[Any], doc_instruction=None):115 doc_instruction = doc_instruction or DOC_INSTRUCTION116 content = docs_to_context_content(docs)117 context = doc_instruction.strip() + "\n" + DOC_TEMPLATE.format(content=content)118 return context119 120 121def maybe_get_doc_context(message, file_input, rag_num_docs: Optional[int] = 3):122 doc_context = None123 if file_input is not None:124 if file_input == RAG_CURRENT_FILE:125 # reuse126 vectorstore = RAG_CURRENT_VECTORSTORE127 print(f'Reuse vectorstore: {file_input}')128 else:129 vectorstore = load_document_split_vectorstore(file_input)130 print(f'New vectorstore: {RAG_CURRENT_FILE} {file_input}')131 RAG_CURRENT_FILE = file_input132 docs = vectorstore.similarity_search(message, k=rag_num_docs)133 doc_context = docs_to_rag_context(docs)134 return doc_context135 136 137def chat_response_stream_multiturn_doc_engine(138 message: str, 139 history: List[Tuple[str, str]], 140 file_input: Optional[str] = None,141 temperature: float = 0.7, 142 max_tokens: int = 1024, 143 system_prompt: Optional[str] = SYSTEM_PROMPT,144 rag_num_docs: Optional[int] = 3,145 doc_instruction: Optional[str] = DOC_INSTRUCTION,146 # profile: Optional[gr.OAuthProfile] = None,147):148 global MODEL_ENGINE, RAG_CURRENT_FILE, RAG_EMBED, RAG_CURRENT_VECTORSTORE149 if len(message) == 0:150 raise gr.Error("The message cannot be empty!")151 152 rag_num_docs = int(rag_num_docs)153 doc_instruction = doc_instruction or DOC_INSTRUCTION154 doc_context = None155 if file_input is not None:156 if file_input == RAG_CURRENT_FILE:157 # reuse158 vectorstore = RAG_CURRENT_VECTORSTORE159 print(f'Reuse vectorstore: {file_input}')160 else:161 vectorstore = load_document_split_vectorstore(file_input)162 print(f'New vectorstore: {RAG_CURRENT_FILE} {file_input}')163 RAG_CURRENT_FILE = file_input164 docs = vectorstore.similarity_search(message, k=rag_num_docs)165 # doc_context = docs_to_rag_context(docs)166 rag_content = docs_to_context_content(docs)167 doc_context = doc_instruction.strip() + "\n" + DOC_TEMPLATE.format(content=rag_content)168 169 if doc_context is not None:170 message = f"{doc_context}\n\n{message}"171 172 for response, num_tokens in chat_response_stream_multiturn_engine(173 message, history, temperature, max_tokens, system_prompt174 ):175 # ! yield another content which is doc_context176 yield response, num_tokens, doc_context177 178 179 180class RagChatInterface(CustomizedChatInterface):181 def __init__(182 self, 183 fn: Callable[..., Any], 184 *, 185 chatbot: gr.Chatbot | None = None, 186 textbox: gr.Textbox | None = None, 187 additional_inputs: str | Component | list[str | Component] | None = None, 188 additional_inputs_accordion_name: str | None = None, 189 additional_inputs_accordion: str | gr.Accordion | None = None, 190 render_additional_inputs_fn: Callable | None = None,191 examples: list[str] | None = None, 192 cache_examples: bool | None = None, 193 title: str | None = None, 194 description: str | None = None, 195 theme: Theme | str | None = None, 196 css: str | None = None, 197 js: str | None = None, 198 head: str | None = None, 199 analytics_enabled: bool | None = None, 200 submit_btn: str | Button | None = "Submit", 201 stop_btn: str | Button | None = "Stop", 202 retry_btn: str | Button | None = "🔄 Retry", 203 undo_btn: str | Button | None = "↩️ Undo", 204 clear_btn: str | Button | None = "🗑️ Clear", 205 autofocus: bool = True, 206 concurrency_limit: int | Literal['default'] | None = "default", 207 fill_height: bool = True208 ):209 try:210 super(gr.ChatInterface, self).__init__(211 analytics_enabled=analytics_enabled,212 mode="chat_interface",213 css=css,214 title=title or "Gradio",215 theme=theme,216 js=js,217 head=head,218 fill_height=fill_height,219 )220 except Exception as e:221 # Handling some old gradio version with out fill_height222 super(gr.ChatInterface, self).__init__(223 analytics_enabled=analytics_enabled,224 mode="chat_interface",225 css=css,226 title=title or "Gradio",227 theme=theme,228 js=js,229 head=head,230 # fill_height=fill_height,231 )232 self.concurrency_limit = concurrency_limit233 self.fn = fn234 self.render_additional_inputs_fn = render_additional_inputs_fn235 self.is_async = inspect.iscoroutinefunction(236 self.fn237 ) or inspect.isasyncgenfunction(self.fn)238 self.is_generator = inspect.isgeneratorfunction(239 self.fn240 ) or inspect.isasyncgenfunction(self.fn)241 self.examples = examples242 if self.space_id and cache_examples is None:243 self.cache_examples = True244 else:245 self.cache_examples = cache_examples or False246 self.buttons: list[Button | None] = []247 248 if additional_inputs:249 if not isinstance(additional_inputs, list):250 additional_inputs = [additional_inputs]251 self.additional_inputs = [252 get_component_instance(i)253 for i in additional_inputs # type: ignore254 ]255 else:256 self.additional_inputs = []257 if additional_inputs_accordion_name is not None:258 print(259 "The `additional_inputs_accordion_name` parameter is deprecated and will be removed in a future version of Gradio. Use the `additional_inputs_accordion` parameter instead."260 )261 self.additional_inputs_accordion_params = {262 "label": additional_inputs_accordion_name263 }264 if additional_inputs_accordion is None:265 self.additional_inputs_accordion_params = {266 "label": "Additional Inputs",267 "open": False,268 }269 elif isinstance(additional_inputs_accordion, str):270 self.additional_inputs_accordion_params = {271 "label": additional_inputs_accordion272 }273 elif isinstance(additional_inputs_accordion, Accordion):274 self.additional_inputs_accordion_params = (275 additional_inputs_accordion.recover_kwargs(276 additional_inputs_accordion.get_config()277 )278 )279 else:280 raise ValueError(281 f"The `additional_inputs_accordion` parameter must be a string or gr.Accordion, not {type(additional_inputs_accordion)}"282 )283 284 with self:285 if title:286 Markdown(287 f"<h1 style='text-align: center; margin-bottom: 1rem'>{self.title}</h1>"288 )289 if description:290 Markdown(description)291 292 if chatbot:293 self.chatbot = chatbot.render()294 else:295 self.chatbot = Chatbot(296 label="Chatbot", scale=1, height=200 if fill_height else None297 )298 299 with Row():300 for btn in [retry_btn, undo_btn, clear_btn]:301 if btn is not None:302 if isinstance(btn, Button):303 btn.render()304 elif isinstance(btn, str):305 btn = Button(btn, variant="secondary", size="sm")306 else:307 raise ValueError(308 f"All the _btn parameters must be a gr.Button, string, or None, not {type(btn)}"309 )310 self.buttons.append(btn) # type: ignore311 312 with Group():313 with Row():314 if textbox:315 textbox.container = False316 textbox.show_label = False317 textbox_ = textbox.render()318 assert isinstance(textbox_, Textbox)319 self.textbox = textbox_320 else:321 self.textbox = Textbox(322 container=False,323 show_label=False,324 label="Message",325 placeholder="Type a message...",326 scale=7,327 autofocus=autofocus,328 )329 if submit_btn is not None:330 if isinstance(submit_btn, Button):331 submit_btn.render()332 elif isinstance(submit_btn, str):333 submit_btn = Button(334 submit_btn,335 variant="primary",336 scale=2,337 min_width=150,338 )339 else:340 raise ValueError(341 f"The submit_btn parameter must be a gr.Button, string, or None, not {type(submit_btn)}"342 )343 if stop_btn is not None:344 if isinstance(stop_btn, Button):345 stop_btn.visible = False346 stop_btn.render()347 elif isinstance(stop_btn, str):348 stop_btn = Button(349 stop_btn,350 variant="stop",351 visible=False,352 scale=2,353 min_width=150,354 )355 else:356 raise ValueError(357 f"The stop_btn parameter must be a gr.Button, string, or None, not {type(stop_btn)}"358 )359 self.num_tokens = Textbox(360 container=False,361 label="num_tokens",362 placeholder="0 tokens",363 scale=1,364 interactive=False,365 # autofocus=autofocus,366 min_width=10367 )368 self.buttons.extend([submit_btn, stop_btn]) # type: ignore369 370 self.fake_api_btn = Button("Fake API", visible=False)371 self.fake_response_textbox = Textbox(label="Response", visible=False)372 (373 self.retry_btn,374 self.undo_btn,375 self.clear_btn,376 self.submit_btn,377 self.stop_btn,378 ) = self.buttons379 380 if examples:381 if self.is_generator:382 examples_fn = self._examples_stream_fn383 else:384 examples_fn = self._examples_fn385 386 self.examples_handler = Examples(387 examples=examples,388 inputs=[self.textbox] + self.additional_inputs,389 outputs=self.chatbot,390 fn=examples_fn,391 )392 393 any_unrendered_inputs = any(394 not inp.is_rendered for inp in self.additional_inputs395 )396 if self.additional_inputs and any_unrendered_inputs:397 with Accordion(**self.additional_inputs_accordion_params): # type: ignore398 if self.render_additional_inputs_fn is not None:399 self.render_additional_inputs_fn()400 else:401 for input_component in self.additional_inputs:402 if not input_component.is_rendered:403 input_component.render()404 405 self.rag_content = gr.Textbox(406 scale=4,407 lines=16,408 label='Retrieved RAG context',409 placeholder="Rag context and instrution will show up here",410 interactive=False411 )412 413 # The example caching must happen after the input components have rendered414 if cache_examples:415 client_utils.synchronize_async(self.examples_handler.cache)416 417 self.saved_input = State()418 self.chatbot_state = (419 State(self.chatbot.value) if self.chatbot.value else State([])420 )421 422 self._setup_events()423 self._setup_api()424 425 def _setup_events(self) -> None:426 from gradio.components import State427 has_on = False428 try:429 from gradio.events import Dependency, EventListenerMethod, on430 has_on = True431 except ImportError as ie:432 has_on = False433 submit_fn = self._stream_fn if self.is_generator else self._submit_fn434 if not self.is_generator:435 raise NotImplementedError(f'should use generator')436 437 if has_on:438 # new version439 submit_triggers = (440 [self.textbox.submit, self.submit_btn.click]441 if self.submit_btn442 else [self.textbox.submit]443 )444 submit_event = (445 on(446 submit_triggers,447 self._clear_and_save_textbox,448 [self.textbox],449 [self.textbox, self.saved_input],450 api_name=False,451 queue=False,452 )453 .then(454 self._display_input,455 [self.saved_input, self.chatbot_state],456 [self.chatbot, self.chatbot_state],457 api_name=False,458 queue=False,459 )460 .then(461 submit_fn,462 [self.saved_input, self.chatbot_state] + self.additional_inputs,463 [self.chatbot, self.chatbot_state, self.num_tokens, self.rag_content],464 api_name=False,465 )466 )467 self._setup_stop_events(submit_triggers, submit_event)468 else:469 raise ValueError(f'Better install new gradio version than 3.44.0')470 471 if self.retry_btn:472 retry_event = (473 self.retry_btn.click(474 self._delete_prev_fn,475 [self.chatbot_state],476 [self.chatbot, self.saved_input, self.chatbot_state],477 api_name=False,478 queue=False,479 )480 .then(481 self._display_input,482 [self.saved_input, self.chatbot_state],483 [self.chatbot, self.chatbot_state],484 api_name=False,485 queue=False,486 )487 .then(488 submit_fn,489 [self.saved_input, self.chatbot_state] + self.additional_inputs,490 [self.chatbot, self.chatbot_state, self.num_tokens, self.rag_content],491 api_name=False,492 )493 )494 self._setup_stop_events([self.retry_btn.click], retry_event)495 496 if self.undo_btn:497 self.undo_btn.click(498 self._delete_prev_fn,499 [self.chatbot_state],500 [self.chatbot, self.saved_input, self.chatbot_state],501 api_name=False,502 queue=False,503 ).then(504 lambda x: x,505 [self.saved_input],506 [self.textbox],507 api_name=False,508 queue=False,509 )510 # Reconfigure clear_btn to stop and clear text box511 512 async def _stream_fn(513 self,514 message: str,515 history_with_input,516 request: Request,517 *args,518 ) -> AsyncGenerator:519 history = history_with_input[:-1]520 inputs, _, _ = special_args(521 self.fn, inputs=[message, history, *args], request=request522 )523 524 if self.is_async:525 generator = self.fn(*inputs)526 else:527 generator = await anyio.to_thread.run_sync(528 self.fn, *inputs, limiter=self.limiter529 )530 generator = SyncToAsyncIterator(generator, self.limiter)531 532 # ! In case of error, yield the previous history & undo any generation before raising error533 try:534 first_response_pack = await async_iteration(generator)535 if isinstance(first_response_pack, (tuple, list)):536 first_response, num_tokens, rag_content = first_response_pack537 else:538 first_response, num_tokens, rag_content = first_response_pack, -1, ""539 update = history + [[message, first_response]]540 yield update, update, f"{num_tokens} toks", rag_content541 except StopIteration:542 update = history + [[message, None]]543 yield update, update, "NaN toks", ""544 except Exception as e:545 yield history, history, "NaN toks", ""546 raise e547 548 try:549 async for response_pack in generator:550 if isinstance(response_pack, (tuple, list)):551 response, num_tokens, rag_content = response_pack552 else:553 response, num_tokens, rag_content = response_pack, "NaN toks", ""554 update = history + [[message, response]]555 yield update, update, f"{num_tokens} toks", rag_content556 except Exception as e:557 yield history, history, "NaN toks", ""558 raise e559 560 561 562@register_demo563class RagChatInterfaceDemo(ChatInterfaceDemo):564 565 @property566 def examples(self):567 return [568 ["Explain how attention works.", "assets/attention_all_you_need.pdf"],569 ["Explain why the sky is blue.", None],570 ]571 572 @property573 def tab_name(self):574 return "RAG Chat"575 576 def create_demo(577 self, 578 title: str | None = None, 579 description: str | None = None, 580 **kwargs581 ) -> gr.Blocks:582 load_embeddings()583 global RAG_EMBED584 # assert RAG_EMBED is not None585 print(F'{RAG_EMBED=}')586 system_prompt = kwargs.get("system_prompt", SYSTEM_PROMPT)587 max_tokens = kwargs.get("max_tokens", MAX_TOKENS)588 temperature = kwargs.get("temperature", TEMPERATURE)589 model_name = kwargs.get("model_name", MODEL_NAME)590 rag_num_docs = kwargs.get("rag_num_docs", 3)591 592 from ..configs import RAG_EMBED_MODEL_NAME593 594 description = (595 description or 596 f"""Upload a long document to ask question with RAG. Check the retrieved RAG text segment below. 597Control `RAG instruction` param to fit your language. Embedding model {RAG_EMBED_MODEL_NAME}."""598 )599 600 additional_inputs = [601 gr.File(label='Upload Document', file_count='single', file_types=['pdf', 'docx', 'txt']),602 gr.Number(value=temperature, label='Temperature', min_width=20), 603 gr.Number(value=max_tokens, label='Max tokens', min_width=20), 604 gr.Textbox(value=system_prompt, label='System prompt', lines=2),605 gr.Number(value=rag_num_docs, label='RAG Top-K', min_width=20),606 gr.Textbox(value=DOC_INSTRUCTION, label='RAG instruction'),607 ]608 def render_additional_inputs_fn():609 additional_inputs[0].render()610 with Row():611 additional_inputs[1].render()612 additional_inputs[2].render()613 additional_inputs[4].render()614 additional_inputs[3].render()615 additional_inputs[5].render()616 617 demo_chat = RagChatInterface(618 chat_response_stream_multiturn_doc_engine,619 chatbot=gr.Chatbot(620 label=model_name,621 bubble_full_width=False,622 latex_delimiters=[623 { "left": "$", "right": "$", "display": False},624 { "left": "$$", "right": "$$", "display": True},625 ],626 show_copy_button=True,627 ),628 textbox=gr.Textbox(placeholder='Type message', lines=1, max_lines=128, min_width=200, scale=8),629 submit_btn=gr.Button(value='Submit', variant="primary", scale=0),630 # ! consider preventing the stop button631 # stop_btn=None,632 title=title,633 description=description,634 additional_inputs=additional_inputs, 635 render_additional_inputs_fn=render_additional_inputs_fn,636 additional_inputs_accordion=gr.Accordion("Additional Inputs", open=True),637 examples=self.examples,638 cache_examples=False,639 )640 return demo_chat641 642 643