Update app.py
Browse files
app.py
CHANGED
@@ -1,117 +1,118 @@
|
|
1 |
-
import
|
2 |
-
import
|
3 |
-
import
|
4 |
-
|
5 |
-
from llama_index.
|
6 |
-
from llama_index.
|
7 |
-
from llama_index.
|
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 |
-
|
|
|
|
1 |
+
import os
|
2 |
+
import gradio as gr
|
3 |
+
import openai
|
4 |
+
import pymongo
|
5 |
+
from llama_index.embeddings.openai import OpenAIEmbedding
|
6 |
+
from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch
|
7 |
+
from llama_index.core import VectorStoreIndex
|
8 |
+
from llama_index.llms.openai import OpenAI
|
9 |
+
|
10 |
+
MONGO_URI = os.environ["MONGO_URI"]
|
11 |
+
|
12 |
+
# MongoDB Atlas Vector Store
|
13 |
+
mongodb_client = pymongo.MongoClient(MONGO_URI)
|
14 |
+
store = MongoDBAtlasVectorSearch(
|
15 |
+
mongodb_client=mongodb_client,
|
16 |
+
db_name="oppenheimer",
|
17 |
+
collection_name="oppenheimer_wiki_chunks",
|
18 |
+
index_name="vector_index",
|
19 |
+
embedding_key="embedding",
|
20 |
+
)
|
21 |
+
|
22 |
+
|
23 |
+
def is_valid_openai_api_key(api_key):
|
24 |
+
client = openai.OpenAI(api_key=api_key)
|
25 |
+
try:
|
26 |
+
client.models.list()
|
27 |
+
except openai.AuthenticationError:
|
28 |
+
return False
|
29 |
+
else:
|
30 |
+
return True
|
31 |
+
|
32 |
+
|
33 |
+
def prepare_query_engine(api_key):
|
34 |
+
# OpenAI Embeddings
|
35 |
+
embed_model = OpenAIEmbedding(
|
36 |
+
model="text-embedding-3-small",
|
37 |
+
embed_batch_size=16,
|
38 |
+
api_key=api_key,
|
39 |
+
max_retries=2,
|
40 |
+
)
|
41 |
+
|
42 |
+
# Loading Index
|
43 |
+
index_loaded = VectorStoreIndex.from_vector_store(
|
44 |
+
vector_store=store, embed_model=embed_model
|
45 |
+
)
|
46 |
+
|
47 |
+
# GPT 3.5 Turbo
|
48 |
+
llm = OpenAI(
|
49 |
+
model="gpt-3.5-turbo-0125", temperature=0, max_tokens=512, api_key=api_key
|
50 |
+
)
|
51 |
+
|
52 |
+
# Query Engine
|
53 |
+
query_engine = index_loaded.as_query_engine(
|
54 |
+
llm=llm, streaming=True, similarity_top_k=3
|
55 |
+
)
|
56 |
+
|
57 |
+
return query_engine
|
58 |
+
|
59 |
+
|
60 |
+
# Generates response using the question answering chain defined earlier
|
61 |
+
def generate(query, api_key):
|
62 |
+
if api_key.strip() == "" or not is_valid_openai_api_key(api_key):
|
63 |
+
yield "Please enter a valid openai api key"
|
64 |
+
else:
|
65 |
+
query_engine = prepare_query_engine(api_key)
|
66 |
+
response = ""
|
67 |
+
try:
|
68 |
+
streaming_response = query_engine.query(query)
|
69 |
+
for token in streaming_response.response_gen:
|
70 |
+
response += token
|
71 |
+
yield response
|
72 |
+
except openai.RateLimitError as rl:
|
73 |
+
yield "RateLimitError - " + str(rl)
|
74 |
+
except Exception as e:
|
75 |
+
yield str(e)
|
76 |
+
|
77 |
+
|
78 |
+
with gr.Blocks() as demo:
|
79 |
+
gr.Markdown(
|
80 |
+
"""
|
81 |
+
# Retrieval Augmented Generation with GPT 3.5 Turbo and MongoDB Atlas Vector Search: Question Answering demo
|
82 |
+
### This demo uses the GPT 3.5 Turbo LLM and MongoDB Atlas Vector Search for fast and performant Retrieval Augmented Generation (RAG).
|
83 |
+
### The context is the new Oppenheimer movie's entire wikipedia page. The movie came out very recently in July, 2023, so the GPT 3.5 turbo model is not aware of it.
|
84 |
+
Retrieval Augmented Generation (RAG) enables us to retrieve just the few small chunks of the document that are relevant to the our query and inject it into our prompt.
|
85 |
+
The model is then able to answer questions by incorporating knowledge from the newly provided document. RAG can be used with thousands of documents, but this demo is limited to just one txt file.
|
86 |
+
"""
|
87 |
+
)
|
88 |
+
OPENAI_API_KEY = gr.Textbox(
|
89 |
+
label="OPENAI_API_KEY",
|
90 |
+
placeholder="Enter your OPENAI_API_KEY",
|
91 |
+
lines=1,
|
92 |
+
type="password",
|
93 |
+
)
|
94 |
+
gr.Markdown("## Enter your question")
|
95 |
+
with gr.Row():
|
96 |
+
with gr.Column():
|
97 |
+
ques = gr.Textbox(label="Question", placeholder="Enter text here", lines=2)
|
98 |
+
with gr.Column():
|
99 |
+
ans = gr.Textbox(label="Answer", lines=4, interactive=False)
|
100 |
+
with gr.Row():
|
101 |
+
with gr.Column():
|
102 |
+
btn = gr.Button("Submit")
|
103 |
+
with gr.Column():
|
104 |
+
clear = gr.ClearButton([ques, ans])
|
105 |
+
|
106 |
+
btn.click(fn=generate, inputs=[ques, OPENAI_API_KEY], outputs=[ans])
|
107 |
+
examples = gr.Examples(
|
108 |
+
examples=[
|
109 |
+
"Who portrayed J. Robert Oppenheimer in the new Oppenheimer movie?",
|
110 |
+
"In the plot of the movie, why did Lewis Strauss resent Robert Oppenheimer?",
|
111 |
+
"What happened while Oppenheimer was a student at the University of Cambridge?",
|
112 |
+
"How much money did the Oppenheimer movie make at the US and global box office?",
|
113 |
+
"What score did the Oppenheimer movie get on Rotten Tomatoes and Metacritic?",
|
114 |
+
],
|
115 |
+
inputs=[ques],
|
116 |
+
)
|
117 |
+
|
118 |
+
demo.queue().launch(debug=True)
|