File size: 4,256 Bytes
9643fb1
 
 
 
24fad8b
9643fb1
 
 
 
 
 
 
 
 
 
 
b8f5c23
 
 
 
 
9643fb1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24fad8b
9643fb1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24fad8b
9643fb1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b8f5c23
 
 
 
 
9643fb1
 
 
 
 
 
 
 
b8f5c23
9643fb1
 
 
 
 
 
 
 
 
 
 
b8f5c23
 
9643fb1
 
 
 
 
 
b8f5c23
9643fb1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b8f5c23
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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
import json
import os
from pathlib import Path

import faiss
from langchain.docstore.document import Document
from langchain.docstore.in_memory import InMemoryDocstore
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores.faiss import FAISS

import gradio as gr

from chain import get_chain

STORE_DIR = "store"
YOUTUBE_EMBED_TEMPLATE = """
<iframe width="354"
    height="200"
    src="{source}&cc_load_policy=1"
    title="YouTube video player"
    frameborder="0"
    allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen>
</iframe>"""


def load_store():
    def keys_to_int(x):
        return {int(k): v for k, v in x.items()}

    index_path = list(Path(STORE_DIR).glob("*.faiss"))
    if len(index_path) == 0:
        raise ValueError("No index found in path")

    index_path = index_path[0]
    index_name = index_path.name.split(".")[0]

    with open(os.path.join(STORE_DIR, f"{index_name}_doc_idx.json"), "r") as f:
        index_to_id = json.load(f, object_hook=keys_to_int)

    with open(os.path.join(STORE_DIR, f"{index_name}_docs.json"), "r") as f:
        docs = json.load(f)

    embeddings = OpenAIEmbeddings()
    return FAISS(
        embedding_function=embeddings.embed_query,
        index=faiss.read_index(str(index_path)),
        docstore=InMemoryDocstore(
            {index_to_id[i]: Document(**doc) for i, doc in enumerate(docs.values())}
        ),
        index_to_docstore_id=index_to_id,
    )


def set_openai_api_key(api_key, agent):
    if api_key:
        os.environ["OPENAI_API_KEY"] = api_key
        vstore = load_store()
        qa_chain = get_chain(vstore)
        os.environ["OPENAI_API_KEY"] = ""
        return qa_chain


def _to_embed(link):
    return link.replace("watch?v=", "embed/").replace("&t=", "?start=")


def chat(inp, history, agent):
    history = history or []
    if agent is None:
        history.append((inp, "Please paste your OpenAI key"))
        return history, history
    output = agent({"question": inp, "chat_history": history})
    answer = output["answer"]
    history.append((inp, answer))
    source_iframes = []
    for source in output["sources"]:
        if "youtube.com" in source:
            source_iframes.append(
                YOUTUBE_EMBED_TEMPLATE.format(source=_to_embed(source))
            )
    source_html = f"""<div style='min-height:200px;display:flex;align-items:center;justify-content:space-around;'>
        {''.join(source_iframes)}
    </div>"""
    return history, history, source_html


with gr.Blocks(css=".gradio-container {background-color: lightgray}") as demo:

    state = gr.State()
    agent_state = gr.State()

    gr.Markdown("<h3><center>ToKBot🤖 - Ask ToKCast Questions</center></h3>")
    openai_api_key_textbox = gr.Textbox(
        placeholder="Paste your OpenAI API key (sk-...)",
        show_label=False,
        lines=1,
        type="password",
    )

    sources = gr.HTML()
    chatbot = gr.Chatbot()
    with gr.Row():
        message = gr.Textbox(
            label="What's your question?",
            placeholder="Type your question here...",
            lines=1,
        )
        submit = gr.Button(value="Send", variant="secondary").style(full_width=False)

    gr.Examples(
        examples=[
            'What does "the beginning of infinity" refer to?',
            "How do memes differ from genes in their replication?",
            "What is the nature of knowledge and how does it grow?",
        ],
        inputs=message,
    )

    gr.HTML(
        """<p>A GPT-3/LangChain bot that answers questions about the TokCast podcast provides relevant video excerpts</p>"""
    )

    gr.HTML(
        "<center>Powered by <a href='https://github.com/hwchase17/langchain'>LangChain 🦜️🔗</a></center>"
    )

    submit.click(
        chat,
        inputs=[message, state, agent_state],
        outputs=[chatbot, state, sources],
    )
    message.submit(
        chat,
        inputs=[message, state, agent_state],
        outputs=[chatbot, state, sources],
    )

    openai_api_key_textbox.change(
        set_openai_api_key,
        inputs=[openai_api_key_textbox, agent_state],
        outputs=[agent_state],
    )

demo.launch()