File size: 4,199 Bytes
7e90a00 84e4090 7e90a00 |
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 |
import os
import gradio as gr
import openai
import pymongo
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch
from llama_index.core import VectorStoreIndex
from llama_index.llms.openai import OpenAI
MONGO_URI = os.environ["MONGO_URI"]
# MongoDB Atlas Vector Store
mongodb_client = pymongo.MongoClient(MONGO_URI)
store = MongoDBAtlasVectorSearch(
mongodb_client=mongodb_client,
db_name="oppenheimer",
collection_name="oppenheimer_wiki_chunks",
index_name="vector_index",
embedding_key="embedding",
)
def is_valid_openai_api_key(api_key):
client = openai.OpenAI(api_key=api_key)
try:
client.models.list()
except openai.AuthenticationError:
return False
else:
return True
def prepare_query_engine(api_key):
# OpenAI Embeddings
embed_model = OpenAIEmbedding(
model="text-embedding-3-small",
embed_batch_size=16,
api_key=api_key,
max_retries=2,
)
# Loading Index
index_loaded = VectorStoreIndex.from_vector_store(
vector_store=store, embed_model=embed_model
)
# GPT 3.5 Turbo
llm = OpenAI(
model="gpt-3.5-turbo-0125", temperature=0, max_tokens=512, api_key=api_key
)
# Query Engine
query_engine = index_loaded.as_query_engine(
llm=llm, streaming=True, similarity_top_k=3
)
return query_engine
# Generates response using the question answering chain defined earlier
def generate(query, api_key):
if api_key.strip() == "" or not is_valid_openai_api_key(api_key):
yield "Please enter a valid openai api key"
else:
query_engine = prepare_query_engine(api_key)
response = ""
try:
streaming_response = query_engine.query(query)
for token in streaming_response.response_gen:
response += token
yield response
except openai.RateLimitError as rl:
yield "RateLimitError - " + str(rl)
except Exception as e:
yield str(e)
with gr.Blocks() as demo:
gr.Markdown(
"""
# Retrieval Augmented Generation with GPT 3.5 Turbo, MongoDB Atlas Vector Search, and LlamaIndex: Question Answering demo
### This demo uses the GPT 3.5 Turbo LLM and MongoDB Atlas Vector Search for fast and performant Retrieval Augmented Generation (RAG).
### 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.
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.
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.
"""
)
OPENAI_API_KEY = gr.Textbox(
label="OPENAI_API_KEY",
placeholder="Enter your OPENAI_API_KEY",
lines=1,
type="password",
)
gr.Markdown("## Enter your question")
with gr.Row():
with gr.Column():
ques = gr.Textbox(label="Question", placeholder="Enter text here", lines=2)
with gr.Column():
ans = gr.Textbox(label="Answer", lines=4, interactive=False)
with gr.Row():
with gr.Column():
btn = gr.Button("Submit")
with gr.Column():
clear = gr.ClearButton([ques, ans])
btn.click(fn=generate, inputs=[ques, OPENAI_API_KEY], outputs=[ans])
examples = gr.Examples(
examples=[
"Who portrayed J. Robert Oppenheimer in the new Oppenheimer movie?",
"In the plot of the movie, why did Lewis Strauss resent Robert Oppenheimer?",
"What happened while Oppenheimer was a student at the University of Cambridge?",
"How much money did the Oppenheimer movie make at the US and global box office?",
"What score did the Oppenheimer movie get on Rotten Tomatoes and Metacritic?",
],
inputs=[ques],
)
demo.queue().launch(debug=True)
|