Update app.py

app.py

@@ -1,8 +1,8 @@
 import streamlit as st
 import torch
 import os
-import time
 import tempfile
+import time
 from threading import Thread
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 from langchain_community.document_loaders import PyPDFLoader, TextLoader
@@ -13,6 +13,10 @@ from langchain.retrievers import BM25Retriever, EnsembleRetriever
 from langchain.schema import Document
 from langchain.docstore.document import Document as LangchainDocument
 
+# --- Avatars ---
+USER_AVATAR = "👤"
+BOT_AVATAR = "🤖"
+
 # --- HF Token ---
 HF_TOKEN = st.secrets["HF_TOKEN"]
 
@@ -31,7 +35,7 @@ with st.sidebar:
 if "messages" not in st.session_state or clear_chat:
     st.session_state.messages = []
 
-# --- Load Model + Tokenizer ---
+# --- Load Model ---
 @st.cache_resource
 def load_model():
     model_id = "tiiuae/falcon-7b-instruct"
@@ -41,7 +45,7 @@ def load_model():
 
 tokenizer, model = load_model()
 
-# --- Process Documents ---
+# --- Load & Chunk Documents ---
 def process_documents(files):
     documents = []
     for file in files:
@@ -49,75 +53,77 @@ def process_documents(files):
         with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
             tmp_file.write(file.read())
             tmp_file_path = tmp_file.name
-
-        if suffix == ".pdf":
-            loader = PyPDFLoader(tmp_file_path)
-        else:
-            loader = TextLoader(tmp_file_path)
-        docs = loader.load()
-        documents.extend(docs)
+        loader = PyPDFLoader(tmp_file_path) if suffix == ".pdf" else TextLoader(tmp_file_path)
+        documents.extend(loader.load())
     return documents
 
 def chunk_documents(documents):
     splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
     return splitter.split_documents(documents)
 
-# --- Build Hybrid Retriever ---
 def build_retrievers(chunks):
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     faiss_vectorstore = FAISS.from_documents(chunks, embeddings)
     faiss_retriever = faiss_vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 5})
     bm25_retriever = BM25Retriever.from_documents([LangchainDocument(page_content=d.page_content) for d in chunks])
     bm25_retriever.k = 5
-    hybrid = EnsembleRetriever(retrievers=[faiss_retriever, bm25_retriever], weights=[0.5, 0.5])
-    return faiss_retriever, hybrid
-
-# --- Inference with Streaming ---
-def generate_stream_response(system_prompt):
+    return faiss_retriever, EnsembleRetriever(retrievers=[faiss_retriever, bm25_retriever], weights=[0.5, 0.5])
+
+# --- Prompt Builder ---
+def build_prompt(history, context=""):
+    conversation = ""
+    for turn in history:
+        role = "User" if turn["role"] == "user" else "Assistant"
+        conversation += f"{role}: {turn['content']}\n"
+    return (
+        "You are DigiTwin, an expert advisor in asset integrity, reliability, inspection, and maintenance "
+        "of topside piping, structural, mechanical systems, floating units, pressure vessels (VII), and pressure safety devices (PSD's).\n\n"
+        f"Context:\n{context}\n\n"
+        f"{conversation}Assistant:"
+    )
+
+# --- Generator ---
+def generate_response(prompt):
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    inputs = tokenizer(system_prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     generation_kwargs = dict(**inputs, streamer=streamer, max_new_tokens=300)
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
-    partial_output = ""
     for token in streamer:
-        partial_output += token
-        yield partial_output
+        yield token
 
-# --- Main App Logic ---
+# --- Main App ---
+retriever = None
 if uploaded_files:
     with st.spinner("Processing documents..."):
         docs = process_documents(uploaded_files)
         chunks = chunk_documents(docs)
-        faiss_retriever, hybrid_retriever = build_retrievers(chunks)
-        retriever = hybrid_retriever if hybrid_toggle else faiss_retriever
-        st.success("Knowledge base ready. Ask your question below.")
-
-    for msg in st.session_state.messages:
-        with st.chat_message(msg["role"]):
-            st.markdown(msg["content"])
-
-    user_input = st.chat_input("💬 Ask DigiTwin something...")
-    if user_input:
-        st.chat_message("user").markdown(user_input)
-        st.session_state.messages.append({"role": "user", "content": user_input})
-
-        with st.chat_message("assistant"):
-            context_docs = retriever.get_relevant_documents(user_input)
-            context_text = "\n".join([doc.page_content for doc in context_docs])
-
-            system_prompt = (
-                "You are DigiTwin, an expert advisor in asset integrity, reliability, inspection, and maintenance "
-                "of topside piping, structural, mechanical systems, floating units, pressure vessels (VII), and pressure safety devices (PSD's).\n\n"
-                f"Context:\n{context_text}\n\n"
-                f"User: {user_input}\nAssistant:"
-            )
-
-            full_response = ""
-            response_area = st.empty()
-            for partial_output in generate_stream_response(system_prompt):
-                full_response = partial_output
-                response_area.markdown(full_response)
-            st.session_state.messages.append({"role": "assistant", "content": full_response})
-else:
-    st.info("👈 Upload one or more PDFs or .txt files to begin.")
+        faiss, hybrid = build_retrievers(chunks)
+        retriever = hybrid if hybrid_toggle else faiss
+        st.success("Documents processed. Ask away!")
+
+for msg in st.session_state.messages:
+    with st.chat_message(msg["role"], avatar=USER_AVATAR if msg["role"] == "user" else BOT_AVATAR):
+        st.markdown(msg["content"])
+
+# --- Chat UI ---
+if prompt := st.chat_input("Ask something based on uploaded documents..."):
+    st.chat_message("user", avatar=USER_AVATAR).markdown(prompt)
+    st.session_state.messages.append({"role": "user", "content": prompt})
+
+    context = ""
+    if retriever:
+        docs = retriever.get_relevant_documents(prompt)
+        context = "\n\n".join([d.page_content for d in docs])
+
+    full_prompt = build_prompt(st.session_state.messages, context=context)
+
+    with st.chat_message("assistant", avatar=BOT_AVATAR):
+        streamer = generate_response(full_prompt)
+        container = st.empty()
+        answer = ""
+        for chunk in streamer:
+            answer += chunk
+            container.markdown(answer + "▌", unsafe_allow_html=True)
+        container.markdown(answer)
+        st.session_state.messages.append({"role": "assistant", "content": answer})