""" Credit to Derek Thomas, derek@huggingface.co """ import os import logging from pathlib import Path from time import perf_counter import gradio as gr from jinja2 import Environment, FileSystemLoader from backend.query_llm import generate_hf, generate_openai from backend.semantic_search import retrieve from backend.reranker import rerank_documents TOP_K = int(os.getenv("TOP_K", 4)) proj_dir = Path(__file__).parent # Setting up the logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Set up the template environment with the templates directory env = Environment(loader=FileSystemLoader(proj_dir / 'templates')) # Load the templates directly from the environment template = env.get_template('template.j2') template_html = env.get_template('template_html.j2') def add_text(history, text): history = [] if history is None else history history = history + [(text, None)] return history, gr.Textbox(value="", interactive=False) def bot(history, chunk_table, embedding_model, llm_model, cross_encoder, top_k_param, rerank_topk ): top_k_param = int(top_k_param) query = history[-1][0] logger.info("bot launched ...") logger.info(f"embedding model: {embedding_model}") logger.info(f"LLM model: {llm_model}") logger.info(f"Cross encoder model: {cross_encoder}") logger.info(f"TopK: {top_k_param}") logger.info(f"ReRank TopK: {rerank_topk}") if not query: raise gr.Warning("Please submit a non-empty string as a prompt") logger.info('Retrieving documents...') # Retrieve documents relevant to query document_start = perf_counter() #documents = retrieve(query, TOP_K) documents = retrieve(query, top_k_param, chunk_table, embedding_model) logger.info(f'Retrived document count: {len(documents)}') if cross_encoder != "None" and len(documents) > 1: documents = rerank_documents(cross_encoder, documents, query, top_k_rerank=rerank_topk) #"cross-encoder/ms-marco-MiniLM-L-6-v2" logger.info(f'ReRank done, document count: {len(documents)}') document_time = perf_counter() - document_start logger.info(f'Finished Retrieving documents in {round(document_time, 2)} seconds...') # Create Prompt prompt = template.render(documents=documents, query=query) prompt_html = template_html.render(documents=documents, query=query) if llm_model == "mistralai/Mistral-7B-Instruct-v0.2": generate_fn = generate_hf if llm_model == "mistralai/Mistral-7B-v0.1": generate_fn = generate_hf if llm_model == "mistralai/Mixtral-8x7B-Instruct-v0.1": generate_fn = generate_hf if llm_model == "gpt-3.5-turbo": generate_fn = generate_openai if llm_model == "gpt-4-turbo-preview": generate_fn = generate_openai #if api_kind == "HuggingFace": # generate_fn = generate_hf #elif api_kind == "OpenAI": # generate_fn = generate_openai #else: # raise gr.Error(f"API {api_kind} is not supported") logger.info(f'Complition started. llm_model: {llm_model}, prompt: {prompt}') history[-1][1] = "" for character in generate_fn(prompt, history[:-1], llm_model): history[-1][1] = character yield history, prompt_html with gr.Blocks() as demo: chatbot = gr.Chatbot( [], elem_id="chatbot", avatar_images=('https://aui.atlassian.com/aui/8.8/docs/images/avatar-person.svg', 'https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.svg'), bubble_full_width=False, show_copy_button=True, show_share_button=True, ) with gr.Row(): txt = gr.Textbox( scale=3, show_label=False, placeholder="Enter text and press enter", container=False, ) txt_btn = gr.Button(value="Submit text", scale=1) #api_kind = gr.Radio(choices=["HuggingFace", # "OpenAI"], value="HuggingFace") chunk_table = gr.Radio(choices=["BGE_CharacterTextSplitter", "BGE_FixedSizeSplitter", "BGE_RecursiveCharacterTextSplitter", "MiniLM_CharacterTextSplitter", "MiniLM_FixedSizeSplitter", "MiniLM_RecursiveCharacterSplitter" ], value="MiniLM_CharacterTextSplitter", label="Chunk table") embedding_model = gr.Radio( choices=[ "BAAI/bge-large-en-v1.5", "sentence-transformers/all-MiniLM-L6-v2", ], value="sentence-transformers/all-MiniLM-L6-v2", label='Embedding model' ) llm_model = gr.Radio( choices=[ "mistralai/Mistral-7B-Instruct-v0.2", "gpt-3.5-turbo", "gpt-4-turbo-preview", "mistralai/Mistral-7B-v0.1", "mistralai/Mixtral-8x7B-Instruct-v0.1" ], value="mistralai/Mistral-7B-Instruct-v0.2", label='LLM' ) cross_encoder = gr.Radio( choices=[ "None", "BAAI/bge-reranker-large", "cross-encoder/ms-marco-MiniLM-L-6-v2", ], value="None", label='Cross-encoder model' ) top_k_param = gr.Radio( choices=[ "5", "10", "20", "50", ], value="5", label='top-K' ) rerank_topk = gr.Radio( choices=[ "5", "10", "20", "50", ], value="5", label='rerank-top-K' ) prompt_html = gr.HTML() # Turn off interactivity while generating if you click txt_msg = txt_btn.click(add_text, [chatbot, txt], [chatbot, txt], queue=False).then( bot, [chatbot, chunk_table, embedding_model, llm_model, cross_encoder, top_k_param, rerank_topk], [chatbot, prompt_html]) # Turn it back on txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False) # Turn off interactivity while generating if you hit enter txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then( bot, [chatbot, chunk_table, embedding_model, llm_model, cross_encoder, top_k_param, rerank_topk], [chatbot, prompt_html]) # Turn it back on txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False) demo.queue() demo.launch(debug=True)