YoutubeQA / app.py
aiswaryamlds's picture
Update app.py
86d2a55
import streamlit as st
from llama_index import VectorStoreIndex, ServiceContext, Document
from llama_index.llms import OpenAI
import openai
from llama_hub.youtube_transcript import YoutubeTranscriptReader
from llama_index import VectorStoreIndex
from langchain.embeddings import HuggingFaceEmbeddings
from llama_index import LangchainEmbedding, ServiceContext
from llama_index.llm_predictor import LLMPredictor
from langchain.llms import LlamaCpp
## For embedding the video, we will use the Hugging Face Sentence Transformers
model_name = "sentence-transformers/all-mpnet-base-v2"
hf = HuggingFaceEmbeddings(
model_name=model_name
)
### We are using LlamaCPP to load the LLAMA-2-18 8 bit quantised model in GGUF format
llm = LlamaCpp(
model_path="codeup-llama-2-13b-chat-hf.Q8_0.gguf",
n_gpu_layers=-1,
n_batch=512,
temperature=0.1,
max_tokens=256,
top_p=1,
verbose=True,
f16_kv=True,
n_ctx=4096,
use_mlock=True,n_threads=4,
stop=["Human:","User:"]
)
## Create a service context object, that will allow us to use the Hugging Face embeddings and llama 2 model as our Language model
llm_predictor=LLMPredictor(llm=llm)
embed_model = LangchainEmbedding(hf)
service_context = ServiceContext.from_defaults(embed_model=embed_model,llm_predictor=llm_predictor)
index=None
### The load data function , takes in youtube_url and allows us to index the youtube video.
def load_data(youtube_url):
print("In Load Data")
if youtube_url.strip()=="":
st.error("Enter A youtube URL")
return None
else:
try:
loader = YoutubeTranscriptReader()
documents = loader.load_data(ytlinks=[youtube_url])
index = VectorStoreIndex.from_documents(documents, service_context=service_context)
return index
except:
print("Enter a valid youtube URL")
st.error("Enter a valid youtube URL")
return None
#### We will have user enter the youtube_url and press submit => which loads the index
index=None
chat_engine=None
### we initiate twp session_state object : clicked and index.
### Clicked: This is set to true when the Submit button is clicked.
### Index: This stores the vector index. By keeping this session state, we allow the index to be persistent till a new yoputube url is enteres
if 'clicked' not in st.session_state:
st.session_state.clicked = False
if 'index' not in st.session_state:
st.session_state.index=None
### click_button-> changes state to Truw when button is clicked
def click_button():
st.session_state.clicked = True
with st.sidebar:
st.title("Youtube QA with Llama 2 Bot")
st.subheader("Upload Documents/URL")
youtube_url = st.sidebar.text_input('Enter Youtube URL', '')
submit_btn=st.sidebar.button('Submit',on_click=click_button)
## When the submit button is clicked, load the data and set the index session_state to the loaded index
if st.session_state.clicked:
print("Going to Load Data")
index=load_data(youtube_url)
st.session_state.index=index
print("Index ",index)
st.session_state.clicked=False # set it to false , so that load_data function is not called for every single user message
#print("Index",index)
print("Index State ",st.session_state.index)
### If the index has been loaded, create the chat_engine object
if st.session_state.index!=None:
chat_engine=st.session_state.index.as_chat_engine(verbose=True,chat_mode="context",service_context=service_context)
print("CHat engine",chat_engine)
if "messages" not in st.session_state.keys():
st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?"}]
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.write(message["content"])
def clear_chat_history():
st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?"}]
st.sidebar.button('Clear Chat History', on_click=clear_chat_history)
if prompt := st.chat_input():
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.write(prompt)
# Generate a new response if last message is not from assistant
if st.session_state.messages[-1]["role"] != "assistant":
full_response = ''
with st.chat_message("assistant"):
with st.spinner("Thinking..."):
print("Calling CHat Engine")
if chat_engine!=None:
response = chat_engine.stream_chat(prompt)
placeholder = st.empty()
for item in response.response_gen:
full_response += item
placeholder.markdown(full_response.strip("Assistant:"))
placeholder.markdown(full_response)
if full_response!="":
message = {"role": "assistant", "content": full_response}
st.session_state.messages.append(message)