File size: 4,060 Bytes
f8817e5 |
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 |
import os
import base64
from io import BytesIO
from PIL import Image
import streamlit as st
from app_config import SYSTEM_PROMPT,MODEL,MAX_TOKENS,TRANSFORMER_MODEL
from langchain.memory import ConversationSummaryBufferMemory
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_groq import ChatGroq
from streamlit_pdf_viewer import pdf_viewer
from pydantic import BaseModel
from langchain.chains import LLMChain
from langchain.prompts import ChatPromptTemplate
from langchain_community.vectorstores import FAISS
from sentence_transformers import SentenceTransformer
from typing import Any
st.title("Hitachi Support Bot")
class Element(BaseModel):
type: str
text: Any
# llm = ChatGoogleGenerativeAI(
# model=MODEL,
# max_tokens=MAX_TOKENS
# )
llm = ChatGroq(model=MODEL,api_key='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')
prompt = ChatPromptTemplate.from_template(SYSTEM_PROMPT)
qa_chain = LLMChain(llm=llm,prompt=prompt)
embeddings = HuggingFaceEmbeddings(model_name=TRANSFORMER_MODEL)
db = FAISS.load_local("faiss_index",embeddings,allow_dangerous_deserialization=True)
st.markdown(
"""
<style>
.st-emotion-cache-janbn0 {
flex-direction: row-reverse;
text-align: right;
}
</style>
""",
unsafe_allow_html=True,
)
def response_generator(question):
relevant_docs = db.similarity_search_with_relevance_scores(question,k=5)
context = ""
relevant_images = []
for d,score in relevant_docs:
if score > 0:
if d.metadata['type'] == 'text':
context += str(d.metadata['original_content'])
elif d.metadata['type'] == 'table':
context += str(d.metadata['original_content'])
elif d.metadata['type'] == 'image':
context += d.page_content
relevant_images.append(d.metadata['original_content'])
result = qa_chain.run({'context':context,"question":question})
return result,relevant_images
with st.sidebar:
st.header("Hitachi Support Bot")
button = st.toggle("View Doc file.")
if button:
pdf_viewer("GPT OUTPUT.pdf")
else:
if "messages" not in st.session_state:
st.session_state.messages=[{"role": "system", "content": SYSTEM_PROMPT}]
if "llm" not in st.session_state:
st.session_state.llm = llm
if "rag_memory" not in st.session_state:
st.session_state.rag_memory = ConversationSummaryBufferMemory(llm=st.session_state.llm, max_token_limit= 5000)
container = st.container(height=700)
for message in st.session_state.messages:
if message["role"] != "system":
if message["role"] == "user":
with container.chat_message(message["role"]):
st.write(message["content"])
if message["role"] == "assistant":
with container.chat_message(message["role"]):
st.write(message["content"])
for i in range(len(message["images"])):
st.image(Image.open(BytesIO(base64.b64decode(message["images"][i].encode('utf-8')))))
if prompt := st.chat_input("Enter your query here... "):
with container.chat_message("user"):
st.write(prompt)
st.session_state.messages.append({"role":"user" , "content":prompt})
with container.chat_message("assistant"):
response,images = response_generator(prompt)
st.write(response)
for i in range(len(images)):
st.markdown("""---""")
st.image(Image.open(BytesIO(base64.b64decode(images[i].encode('utf-8')))))
st.markdown("""---""")
st.session_state.rag_memory.save_context({'input': prompt}, {'output': response})
st.session_state.messages.append({"role":"assistant" , "content":response,'images':images}) |