Update app.py

app.py

@@ -10,21 +10,21 @@ from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
 from langchain.schema import Document
 
-# --- HF Token ---
+# --- Hugging Face Token ---
 HF_TOKEN = st.secrets["HF_TOKEN"]
 
 # --- Page Config ---
 st.set_page_config(page_title="DigiTwin RAG", page_icon="📂", layout="centered")
 st.title("📂 DigiTs the Twin")
 
-# --- Upload Files Sidebar ---
+# --- File Upload UI ---
 with st.sidebar:
     st.header("📄 Upload Knowledge Files")
     uploaded_files = st.file_uploader("Upload PDFs or .txt files", accept_multiple_files=True, type=["pdf", "txt"])
     if uploaded_files:
         st.success(f"{len(uploaded_files)} file(s) uploaded")
 
-# --- Model Loading ---
+# --- Load Model & Tokenizer ---
 @st.cache_resource
 def load_model():
     tokenizer = AutoTokenizer.from_pretrained("amiguel/GM_Qwen1.8B_Finetune", trust_remote_code=True, token=HF_TOKEN)
@@ -39,7 +39,7 @@ def load_model():
 
 model, tokenizer = load_model()
 
-# --- Prompt Helper ---
+# --- System Prompt ---
 SYSTEM_PROMPT = (
     "You are DigiTwin, a digital expert and senior topside engineer specializing in inspection and maintenance "
     "of offshore piping systems, structural elements, mechanical equipment, floating production units, pressure vessels "
@@ -48,7 +48,7 @@ SYSTEM_PROMPT = (
     "field experience, industry regulations, and proven methodologies in asset integrity and reliability engineering."
 )
 
-
+# --- Prompt Builder ---
 def build_prompt(messages, context=""):
     prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}\n\nContext:\n{context}<|im_end|>\n"
     for msg in messages:
@@ -57,17 +57,15 @@ def build_prompt(messages, context=""):
     prompt += "<|im_start|>assistant\n"
     return prompt
 
-
-# --- RAG Embedding and Search ---
+# --- Embed Uploaded Documents ---
 @st.cache_resource
 def embed_uploaded_files(files):
     raw_docs = []
     for f in files:
-        file_path = f"/tmp/{f.name}"
-        with open(file_path, "wb") as out_file:
+        path = f"/tmp/{f.name}"
+        with open(path, "wb") as out_file:
             out_file.write(f.read())
-
-        loader = PyPDFLoader(file_path) if f.name.endswith(".pdf") else TextLoader(file_path)
+        loader = PyPDFLoader(path) if f.name.endswith(".pdf") else TextLoader(path)
         raw_docs.extend(loader.load())
 
     splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=64)
@@ -78,7 +76,7 @@ def embed_uploaded_files(files):
 
 retriever = embed_uploaded_files(uploaded_files) if uploaded_files else None
 
-# --- Streaming Response ---
+# --- Streaming Generator ---
 def generate_response(prompt_text):
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)
@@ -95,37 +93,56 @@ def generate_response(prompt_text):
     thread.start()
     return streamer
 
-# --- Avatars & Messages ---
+# --- Avatars ---
 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"
 
+# --- Initialize Chat Memory ---
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
+# --- Display Message History ---
 for msg in st.session_state.messages:
-    avatar = USER_AVATAR if msg["role"] == "user" else BOT_AVATAR
-    with st.chat_message(msg["role"], avatar=avatar):
+    with st.chat_message(msg["role"], avatar=USER_AVATAR if msg["role"] == "user" else BOT_AVATAR):
         st.markdown(msg["content"])
 
-# --- Chat UI ---
+# --- Chat Interface ---
 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 = ""
+    docs = []
     if retriever:
         docs = retriever.similarity_search(prompt, k=3)
-        context = "\n\n".join([d.page_content for d in docs])
+        context = "\n\n".join([doc.page_content for doc in docs])
 
-    full_prompt = build_prompt(st.session_state.messages, context=context)
+    # Limit to last 6 messages for memory
+    recent_messages = st.session_state.messages[-6:]
+    full_prompt = build_prompt(recent_messages, context)
 
     with st.chat_message("assistant", avatar=BOT_AVATAR):
-        start_time = time.time()
-        streamer = generate_response(full_prompt)
+        start = time.time()
         container = st.empty()
         answer = ""
-        for chunk in streamer:
+
+        for chunk in generate_response(full_prompt):
             answer += chunk
             container.markdown(answer + "▌", unsafe_allow_html=True)
         container.markdown(answer)
-        st.session_state.messages.append({"role": "assistant", "content": answer})
+
+        end = time.time()
+        st.session_state.messages.append({"role": "assistant", "content": answer})
+
+        input_tokens = len(tokenizer(full_prompt)["input_ids"])
+        output_tokens = len(tokenizer(answer)["input_ids"])
+        speed = output_tokens / (end - start)
+
+        with st.expander("📊 Debug Info"):
+            st.caption(
+                f"🔑 Input Tokens: {input_tokens} | Output Tokens: {output_tokens} | "
+                f"🕒 Speed: {speed:.1f} tokens/sec"
+            )
+            for i, doc in enumerate(docs):
+                st.markdown(f"**Chunk #{i+1}**")
+                st.code(doc.page_content.strip()[:500])