ruti3/rag-exercise

1"""2Credit to Derek Thomas, derek@huggingface.co3"""4import os5import logging6from pathlib import Path7from time import perf_counter8 9import gradio as gr10from jinja2 import Environment, FileSystemLoader11 12from backend.query_llm import generate_hf, generate_openai13from backend.semantic_search import retrieve14 15 16TOP_K = int(os.getenv("TOP_K", 4))17 18proj_dir = Path(__file__).parent19# Setting up the logging20logging.basicConfig(level=logging.INFO)21logger = logging.getLogger(__name__)22 23# Set up the template environment with the templates directory24env = Environment(loader=FileSystemLoader(proj_dir / 'templates'))25 26# Load the templates directly from the environment27template = env.get_template('template.j2')28template_html = env.get_template('template_html.j2')29 30 31def add_text(history, text):32    history = [] if history is None else history33    history = history + [(text, None)]34    return history, gr.Textbox(value="", interactive=False)35 36 37def bot(history, api_kind):38    query = history[-1][0]39 40    if not query:41        raise gr.Warning("Please submit a non-empty string as a prompt")42 43    logger.info('Retrieving documents...')44    # Retrieve documents relevant to query45    document_start = perf_counter()46 47    documents = retrieve(query, TOP_K)48 49    document_time = perf_counter() - document_start50    logger.info(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')51 52    # Create Prompt53    prompt = template.render(documents=documents, query=query)54    prompt_html = template_html.render(documents=documents, query=query)55 56    if api_kind == "HuggingFace":57         generate_fn = generate_hf58    elif api_kind == "OpenAI":59         generate_fn = generate_openai60    else:61         raise gr.Error(f"API {api_kind} is not supported")62 63    history[-1][1] = ""64    for character in generate_fn(prompt, history[:-1]):65        history[-1][1] = character66        yield history, prompt_html67 68 69with gr.Blocks() as demo:70    chatbot = gr.Chatbot(71            [],72            elem_id="chatbot",73            avatar_images=('https://aui.atlassian.com/aui/8.8/docs/images/avatar-person.svg',74                           'https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.svg'),75            bubble_full_width=False,76            show_copy_button=True,77            show_share_button=True,78            )79 80    with gr.Row():81        txt = gr.Textbox(82                scale=3,83                show_label=False,84                placeholder="Enter text and press enter",85                container=False,86                )87        txt_btn = gr.Button(value="Submit text", scale=1)88 89    api_kind = gr.Radio(choices=["HuggingFace", "OpenAI"], value="HuggingFace")90 91    prompt_html = gr.HTML()92    # Turn off interactivity while generating if you click93    txt_msg = txt_btn.click(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(94            bot, [chatbot, api_kind], [chatbot, prompt_html])95 96    # Turn it back on97    txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)98 99    # Turn off interactivity while generating if you hit enter100    txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(101            bot, [chatbot, api_kind], [chatbot, prompt_html])102 103    # Turn it back on104    txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)105 106demo.queue()107demo.launch(debug=True)108