captain-awesome's picture
Update app.py
2f9704c verified
raw
history blame
5.4 kB
import streamlit as st
from langchain_core.messages import AIMessage, HumanMessage
from langchain_community.document_loaders import WebBaseLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
# from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_community.llms import CTransformers
from ctransformers import AutoModelForCausalLM
from langchain.llms import HuggingFaceHub
from transformers import AutoModelForCausalLM, AutoTokenizer
import os
# from dotenv import load_dotenv
# load_dotenv()
def get_vector_store_from_url(url):
model_name = "BAAI/bge-large-en"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': False}
embeddings = HuggingFaceBgeEmbeddings(
model_name=model_name,
model_kwargs=model_kwargs,
encode_kwargs=encode_kwargs
)
loader = WebBaseLoader(url)
document = loader.load()
# split the document into chunks
text_splitter = RecursiveCharacterTextSplitter()
document_chunks = text_splitter.split_documents(document)
# create a vectorstore from the chunks
# vector_store = Chroma.from_documents(document_chunks, OpenAIEmbeddings())
vector_store = Chroma.from_documents(document_chunks, embeddings)
return vector_store
def get_context_retriever_chain(vector_store,llm):
# llm = ChatOpenAI()
llm = llm
retriever = vector_store.as_retriever()
prompt = ChatPromptTemplate.from_messages([
MessagesPlaceholder(variable_name="chat_history"),
("user", "{input}"),
("user", "Given the above conversation, generate a search query to look up in order to get information relevant to the conversation")
])
retriever_chain = create_history_aware_retriever(llm, retriever, prompt)
return retriever_chain
def get_conversational_rag_chain(retriever_chain,llm):
llm=llm
prompt = ChatPromptTemplate.from_messages([
("system", "Answer the user's questions based on the below context:\n\n{context}"),
MessagesPlaceholder(variable_name="chat_history"),
("user", "{input}"),
])
stuff_documents_chain = create_stuff_documents_chain(llm,prompt)
return create_retrieval_chain(retriever_chain, stuff_documents_chain)
def get_response(user_input):
# llm = CTransformers(
# # model = "TheBloke/Mistral-7B-Instruct-v0.2-GGUF",
# model= "TheBloke/Llama-2-7B-Chat-GGUF",
# model_file = "llama-2-7b-chat.Q3_K_S.gguf",
# model_type="llama",
# max_new_tokens = 300,
# temperature = 0.3,
# lib="avx2", # for CPU
# )
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
# llm = HuggingFaceHub(
# repo_id=llm_model,
# model_kwargs={"temperature": 0.3, "max_new_tokens": 250, "top_k": 3}
# )
llm = transformers.AutoModelForCausalLM.from_pretrained(
model_name,
trust_remote_code=True,
torch_dtype=torch.bfloat16,
device_map='auto'
)
retriever_chain = get_context_retriever_chain(st.session_state.vector_store,llm)
conversation_rag_chain = get_conversational_rag_chain(retriever_chain,llm)
response = conversation_rag_chain.invoke({
"chat_history": st.session_state.chat_history,
"input": user_query
})
return response['answer']
# app config
st.set_page_config(page_title= "Chat with Websites", page_icon="🤖")
st.title("Chat with Websites")
#sidebar
with st.sidebar:
st.header("Settings")
website_url = st.text_input("Website URL")
# openai_apikey = st.text_input("Enter your OpenAI API key")
if (website_url is None or website_url == ""):
st.info("Please ensure if website URL is entered")
else:
if "chat_history" not in st.session_state:
st.session_state.chat_history = [
AIMessage(content = "Hello, I am a bot. How can I help you"),
]
if "vector_store" not in st.session_state:
st.session_state.vector_store = get_vector_store_from_url(website_url)
#user_input
user_query = st.chat_input("Type your message here...")
if user_query is not None and user_query !="":
response = get_response(user_query)
st.session_state.chat_history.append(HumanMessage(content=user_query))
st.session_state.chat_history.append(AIMessage(content=response))
#conversation
for message in st.session_state.chat_history:
if isinstance(message, AIMessage): # checking if the messsage is the instance of an AI message
with st.chat_message("AI"):
st.write(message.content)
elif isinstance(message, HumanMessage): # checking if the messsage is the instance of a Human
with st.chat_message("Human"):
st.write(message.content)