import logging from functools import partial from pathlib import Path from time import perf_counter import gradio as gr from gradio_rich_textbox import RichTextbox from jinja2 import Environment, FileSystemLoader from transformers import AutoTokenizer from backend.query_llm import check_endpoint_status, generate from backend.semantic_search import retriever 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') # Initialize tokenizer tokenizer = AutoTokenizer.from_pretrained('derek-thomas/jais-13b-chat-hf') # Examples examples = ['من كان طرفي معركة اكتيوم البحرية؟', 'لم السماء زرقاء؟', "من فاز بكأس العالم للرجال في عام 2014؟", ] 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, hyde=False): top_k = 5 query = history[-1][0] logger.warning('Retrieving documents...') # Retrieve documents relevant to query document_start = perf_counter() if hyde: hyde_document = generate(f"Write a wikipedia article intro paragraph to answer this query: {query}").split( '### Response: [|AI|]')[-1] logger.warning(hyde_document) documents = retriever(hyde_document, top_k=top_k) else: documents = retriever(query, top_k=top_k) document_time = perf_counter() - document_start logger.warning(f'Finished Retrieving documents in {round(document_time, 2)} seconds...') # Function to count tokens def count_tokens(text): return len(tokenizer.encode(text)) # Create Prompt prompt = template.render(documents=documents, query=query) # Check if the prompt is too long token_count = count_tokens(prompt) while token_count > 2048: # Shorten your documents here. This is just a placeholder for the logic you'd use. documents.pop() # Remove the last document prompt = template.render(documents=documents, query=query) # Re-render the prompt token_count = count_tokens(prompt) # Re-count tokens prompt_html = template_html.render(documents=documents, query=query) history[-1][1] = "" response = generate(prompt) history[-1][1] = response.split('### Response: [|AI|]')[-1] return history, prompt_html intro_md = """ # Arabic RAG This is a project to demonstrate Retreiver Augmented Generation (RAG) in Arabic and English. It uses [Arabic Wikipedia](https://ar.wikipedia.org/wiki) as a base to answer questions you have. A retriever ([sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2](https://huggingface.co/sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2/discussions/8)) will find the articles relevant to your query and include them in a prompt so the reader ([core42/jais-13b-chat](https://huggingface.co/core42/jais-13b-chat)) can then answer your questions on it. You can see the prompt clearly displayed below the chatbot to understand what is going to the LLM. # Read this if you get an error I'm using [Inference Endpoint's](https://huggingface.co/inference-endpoints) [Scale to Zero](https://huggingface.co/docs/inference-endpoints/main/en/autoscaling#scaling-to-0) to save money on GPUs. If the staus is "scaledToZero" click **Wake Up Endpoint** to wake it up. You will get an `error` and it will take ~4 minutes to wake up. This is expected, if you dont like it please give me a free GPU with enough VRAM. """ def process_example(text, history=[]): history = history + [[text, None]] return bot(history) # hyde_prompt_html = gr.HTML() with gr.Blocks() as demo: gr.Markdown(intro_md) endpoint_status = RichTextbox(check_endpoint_status, label="Inference Endpoint Status", every=1) wakeup_endpoint = gr.Button('Click to Wake Up Endpoint') with gr.Tab("Arabic-RAG"): 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 query in Arabic or English and press enter", container=False, ) txt_btn = gr.Button(value="Submit text", scale=1) # gr.Examples(examples, txt) prompt_html = gr.HTML() gr.Examples( examples=examples, inputs=txt, outputs=[chatbot, prompt_html], fn=process_example, cache_examples=True, ) # prompt_html.render() # Turn off interactivity while generating if you click txt_msg = txt_btn.click(add_text, [chatbot, txt], [chatbot, txt], queue=False).then( bot, chatbot, [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, [chatbot, prompt_html]) # Turn it back on txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False) # Easy to turn this on when I want to # with gr.Tab("Arabic-RAG + HyDE"): # hyde_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(): # hyde_txt = gr.Textbox( # scale=3, # show_label=False, # placeholder="Enter text and press enter", # container=False, # ) # hyde_txt_btn = gr.Button(value="Submit text", scale=1) # # hyde_prompt_html = gr.HTML() # gr.Examples( # examples=examples, # inputs=hyde_txt, # outputs=[hyde_chatbot, hyde_prompt_html], # fn=process_example, # cache_examples=True, ) # # prompt_html.render() # # Turn off interactivity while generating if you click # hyde_txt_msg = hyde_txt_btn.click(add_text, [hyde_chatbot, hyde_txt], [hyde_chatbot, hyde_txt], # queue=False).then( # partial(bot, hyde=True), [hyde_chatbot], [hyde_chatbot, hyde_prompt_html]) # # # Turn it back on # hyde_txt_msg.then(lambda: gr.Textbox(interactive=True), None, [hyde_txt], queue=False) # # # Turn off interactivity while generating if you hit enter # hyde_txt_msg = hyde_txt.submit(add_text, [hyde_chatbot, hyde_txt], [hyde_chatbot, hyde_txt], queue=False).then( # partial(bot, hyde=True), [hyde_chatbot], [hyde_chatbot, hyde_prompt_html]) # # # Turn it back on # hyde_txt_msg.then(lambda: gr.Textbox(interactive=True), None, [hyde_txt], queue=False) wakeup_endpoint.click(partial(generate,'Wakeup')) demo.queue() demo.launch(debug=True)