Update app.py

app.py

@@ -13,23 +13,23 @@ from langchain.schema import Document
 from langchain.docstore.document import Document as LangchainDocument
 
 # --- Avatars ---
-USER_AVATAR = "👤"
-BOT_AVATAR = "🤖"
+USER_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/9904d9a0d445ab0488cf7395cb863cce7621d897/USER_AVATAR.png"
+BOT_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/991f4c6e4e1dc7a8e24876ca5aae5228bcdb4dba/Ataliba_Avatar.jpg"
 
-# --- HF Token ---
+# --- Hugging Face Token ---
 HF_TOKEN = st.secrets["HF_TOKEN"]
 
-# --- Page Config ---
-st.set_page_config(page_title="Hybrid RAG with Streaming", page_icon="📄", layout="centered")
-st.title("📄 Hybrid Search + Streaming Chat")
+# --- Page Setup ---
+st.set_page_config(page_title="Hybrid RAG Chat", page_icon="🤖", layout="centered")
+st.title("🤖 DigiTwin - Hybrid Search + Streaming")
 
 # --- Sidebar Upload ---
 with st.sidebar:
     st.header("📤 Upload Documents")
-    uploaded_files = st.file_uploader("Upload PDFs or .txt files", type=["pdf", "txt"], accept_multiple_files=True)
+    uploaded_files = st.file_uploader("PDFs or .txt files only", type=["pdf", "txt"], accept_multiple_files=True)
     clear_chat = st.button("🧹 Clear Conversation")
 
-# --- Session State ---
+# --- Chat Memory ---
 if "messages" not in st.session_state or clear_chat:
     st.session_state.messages = []
 
@@ -43,87 +43,87 @@ def load_model():
 
 tokenizer, model = load_model()
 
-# --- Load & Chunk Documents ---
+# --- Document Processing ---
 def process_documents(files):
     documents = []
     for file in files:
         suffix = ".pdf" if file.name.endswith(".pdf") else ".txt"
-        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
-            tmp_file.write(file.read())
-            tmp_file_path = tmp_file.name
-        loader = PyPDFLoader(tmp_file_path) if suffix == ".pdf" else TextLoader(tmp_file_path)
+        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+            tmp.write(file.read())
+            path = tmp.name
+        loader = PyPDFLoader(path) if suffix == ".pdf" else TextLoader(path)
         documents.extend(loader.load())
     return documents
 
-def chunk_documents(documents):
+def chunk_documents(docs):
     splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-    return splitter.split_documents(documents)
+    return splitter.split_documents(docs)
 
 def build_hybrid_retriever(chunks):
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-    faiss_store = FAISS.from_documents(chunks, embeddings)
-    faiss_retriever = faiss_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
+    faiss = FAISS.from_documents(chunks, embeddings)
+    faiss_ret = faiss.as_retriever(search_type="similarity", search_kwargs={"k": 5})
+    bm25 = BM25Retriever.from_documents([LangchainDocument(page_content=c.page_content) for c in chunks])
+    bm25.k = 5
+    return EnsembleRetriever(retrievers=[faiss_ret, bm25], weights=[0.5, 0.5])
 
-    bm25_retriever = BM25Retriever.from_documents([LangchainDocument(page_content=d.page_content) for d in chunks])
-    bm25_retriever.k = 5
+# --- Prompt Builder ---
+def build_prompt(history, context=""):
+    dialog = ""
+    for msg in history:
+        role = "User" if msg["role"] == "user" else "Assistant"
+        dialog += f"{role}: {msg['content']}\n"
+    return f"""You are DigiTwin, a highly professional and experienced assistant in inspection, integrity, and maintenance of topside equipment, piping systems, pressure vessels, structures, and safety systems.
 
-    hybrid = EnsembleRetriever(retrievers=[faiss_retriever, bm25_retriever], weights=[0.5, 0.5])
-    return hybrid
+Use the following context to provide expert-level answers.
 
-# --- Prompt Construction ---
-def build_prompt(history, context=""):
-    prompt = (
-        "You are DigiTwin, an expert in reliability, inspection, and maintenance of piping, structures, vessels, and topside assets.\n"
-        f"Use the following context to help answer questions:\n\n{context}\n\n"
-    )
-    for turn in history:
-        role = "User" if turn["role"] == "user" else "Assistant"
-        prompt += f"{role}: {turn['content']}\n"
-    prompt += "Assistant:"
-    return prompt
-
-# --- Generator for Streaming ---
-def generate_streaming_response(prompt):
+Context:
+{context}
+
+{dialog}
+Assistant:"""
+
+# --- Response Generator ---
+def generate_response(prompt):
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    thread = Thread(target=model.generate, kwargs={**inputs, "streamer": streamer, "max_new_tokens": 300})
-    thread.start()
+    Thread(target=model.generate, kwargs={**inputs, "streamer": streamer, "max_new_tokens": 300}).start()
     output = ""
     for token in streamer:
         output += token
         yield output
 
-# --- Run Document Processing and Retrieval ---
+# --- Retrieval Logic ---
 retriever = None
 if uploaded_files:
-    with st.spinner("📚 Processing documents..."):
+    with st.spinner("🔍 Indexing documents..."):
         docs = process_documents(uploaded_files)
         chunks = chunk_documents(docs)
         retriever = build_hybrid_retriever(chunks)
-    st.success("✅ Document processing complete.")
+    st.success("✅ Documents ready for hybrid search.")
 
-# --- Display Past Messages ---
+# --- Display Conversation ---
 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"])
 
-# --- Main Chat Input ---
-if prompt := st.chat_input("Ask a question..."):
-    st.chat_message("user", avatar=USER_AVATAR).markdown(prompt)
-    st.session_state.messages.append({"role": "user", "content": prompt})
+# --- Chat Input ---
+if query := st.chat_input("Ask DigiTwin anything..."):
+    st.chat_message("user", avatar=USER_AVATAR).markdown(query)
+    st.session_state.messages.append({"role": "user", "content": query})
 
     context = ""
     if retriever:
-        docs = retriever.get_relevant_documents(prompt)
+        docs = retriever.get_relevant_documents(query)
         context = "\n\n".join([doc.page_content for doc in docs])
 
-    full_prompt = build_prompt(st.session_state.messages, context=context)
+    full_prompt = build_prompt(st.session_state.messages, context)
 
     with st.chat_message("assistant", avatar=BOT_AVATAR):
-        response_container = st.empty()
+        container = st.empty()
         answer = ""
-        for chunk in generate_streaming_response(full_prompt):
+        for chunk in generate_response(full_prompt):
             answer = chunk
-            response_container.markdown(answer + "▌", unsafe_allow_html=True)
-        response_container.markdown(answer)
+            container.markdown(answer + "▌", unsafe_allow_html=True)
+        container.markdown(answer)
         st.session_state.messages.append({"role": "assistant", "content": answer})