File size: 3,789 Bytes
b7f4e8c
 
 
 
 
 
 
 
 
 
 
 
 
 
93c49cb
 
b7f4e8c
 
93c49cb
b7f4e8c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
93c49cb
b7f4e8c
 
93c49cb
b7f4e8c
 
 
 
93c49cb
b7f4e8c
93c49cb
b7f4e8c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
93c49cb
b7f4e8c
 
93c49cb
b7f4e8c
93c49cb
b7f4e8c
93c49cb
b7f4e8c
 
 
 
 
 
 
 
 
93c49cb
 
 
 
 
 
 
 
b7f4e8c
 
 
93c49cb
 
 
 
 
b7f4e8c
 
 
 
 
 
93c49cb
 
6b2076c
 
 
 
 
 
 
 
b7f4e8c
 
 
93c49cb
 
 
6b2076c
93c49cb
 
 
6b2076c
93c49cb
b7f4e8c
 
 
93c49cb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
"""
Credit to Derek Thomas, [email protected]
"""
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 dotenv import load_dotenv
load_dotenv()

TOP_K = int(os.getenv("TOP_K", 4))
HF_TOKEN = os.getenv("HF_TOKEN")

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):
    logger.info(f'Adding text: {text}')
    history = [] if history is None else history
    history = history + [(text, None)]
    logger.info(f'Updated history: {history}')
    return history, gr.Textbox(value="", interactive=False)


def bot(history, api_kind):
    logger.info(f'Bot function called with history: {history} and api_kind: {api_kind}')
    query = history[-1][0]
    logger.info(f'Query: {query}')

    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)

    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)
    logger.info(f'Prompt created: {prompt}')

    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")

    history[-1][1] = ""
    for character in generate_fn(prompt, history[:-1]):
        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")

    prompt_html = gr.HTML()

    # Turn off interactivity while generating if you click
    txt_msg = txt_btn.click(
            fn=add_text,
            inputs=[chatbot, txt],
            outputs=[chatbot, txt]
        ).then(
            fn=bot,
            inputs=[chatbot, api_kind],
            outputs=[chatbot, prompt_html]
        )
    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(
        fn=add_text,
        inputs=[chatbot, txt],
        outputs=[chatbot, txt]
    ).then(
        fn=bot,
        inputs=[chatbot, api_kind],
        outputs=[chatbot, prompt_html]
    )
    txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)

demo.queue()
demo.launch(debug=True)