|
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 = 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}) |
