Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import torch
 import os
 import tempfile
+import time
 from threading import Thread
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 from langchain_community.document_loaders import PyPDFLoader, TextLoader
@@ -21,12 +22,13 @@ HF_TOKEN = st.secrets["HF_TOKEN"]
 
 # --- Page Setup ---
 st.set_page_config(page_title="Hybrid RAG Chat", page_icon="🤖", layout="centered")
-st.title("🤖 DigiTwin - Hybrid Search + Streaming")
+st.title("🤖 DigiTwin Streaming")
 
 # --- Sidebar Upload ---
 with st.sidebar:
     st.header("📤 Upload Documents")
     uploaded_files = st.file_uploader("PDFs or .txt files only", type=["pdf", "txt"], accept_multiple_files=True)
+    max_tokens = st.slider("🧠 Max Response Tokens", 100, 2048, 512, step=50)
     clear_chat = st.button("🧹 Clear Conversation")
 
 # --- Chat Memory ---
@@ -36,7 +38,7 @@ if "messages" not in st.session_state or clear_chat:
 # --- Load LLM ---
 @st.cache_resource
 def load_model():
-    model_id = "amiguel/GM_Qwen1.8B_Finetune"
+    model_id = "tiiuae/falcon-7b-instruct"
     tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
     model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto", token=HF_TOKEN)
     return tokenizer, model
@@ -73,9 +75,7 @@ def build_prompt(history, context=""):
     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.
-
-Use the following context to provide expert-level answers.
+    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. Use the following context to provide expert-level answers.
 
 Context:
 {context}
@@ -84,10 +84,10 @@ Context:
 Assistant:"""
 
 # --- Response Generator ---
-def generate_response(prompt):
+def generate_response(prompt, max_tokens):
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    Thread(target=model.generate, kwargs={**inputs, "streamer": streamer, "max_new_tokens": 300}).start()
+    Thread(target=model.generate, kwargs={**inputs, "streamer": streamer, "max_new_tokens": max_tokens}).start()
     output = ""
     for token in streamer:
         output += token
@@ -113,17 +113,39 @@ if query := st.chat_input("Ask DigiTwin anything..."):
     st.session_state.messages.append({"role": "user", "content": query})
 
     context = ""
+    matched_chunks = []
     if retriever:
-        docs = retriever.get_relevant_documents(query)
-        context = "\n\n".join([doc.page_content for doc in docs])
+        matched_chunks = retriever.get_relevant_documents(query)
+        context = "\n\n".join([doc.page_content for doc in matched_chunks])
 
     full_prompt = build_prompt(st.session_state.messages, context)
 
     with st.chat_message("assistant", avatar=BOT_AVATAR):
+        start_time = time.time()
         container = st.empty()
         answer = ""
-        for chunk in generate_response(full_prompt):
+
+        for chunk in generate_response(full_prompt, max_tokens):
             answer = chunk
             container.markdown(answer + "▌", unsafe_allow_html=True)
         container.markdown(answer)
+
+        end_time = time.time()
+        input_tokens = len(tokenizer(full_prompt)["input_ids"])
+        output_tokens = len(tokenizer(answer)["input_ids"])
+        speed = output_tokens / (end_time - start_time)
+
         st.session_state.messages.append({"role": "assistant", "content": answer})
+
+        # RAG Debug Info
+        with st.expander("📊 Response Stats & RAG Debug"):
+            st.caption(
+                f"🔑 Input Tokens: {input_tokens} | Output Tokens: {output_tokens} | "
+                f"🕒 Speed: {speed:.1f} tokens/sec"
+            )
+            for i, doc in enumerate(matched_chunks):
+                score = getattr(doc, "score", None)
+                metadata = doc.metadata if hasattr(doc, "metadata") else {}
+                st.markdown(f"**Chunk #{i+1}**")
+                st.code(doc.page_content.strip()[:500])
+                st.text(f"🔍 Similarity Score: {score if score else 'N/A'} | Metadata: {metadata}")