Update app.py

app.py

@@ -1,144 +1,129 @@
 import streamlit as st
+import torch
 import os
-import time
-import PyPDF2
-from docx import Document
-import pandas as pd
-from dotenv import load_dotenv
-from unsloth import FastLanguageModel
-from transformers import AutoTokenizer
-
-# Load environment variables
-load_dotenv()
-
-# Avatars and bios
-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"
-
-ATALIBA_BIO = """
-**I am Ataliba Miguel's Digital Twin** 🤖
-
-**Background:**
-- 🎓 Mechanical Engineering (BSc)
-- ⛽ Oil & Gas Engineering (MSc Specialization)
-- 🔧 17+ years in Oil & Gas Industry
-- 🔍 Current: Topside Inspection Methods Engineer @ TotalEnergies
-- 🤖 AI Practitioner Specialist
-- 🚀 Founder of ValonyLabs (AI solutions for industrial corrosion, retail analytics, and KPI monitoring)
-
-**Capabilities:**
-- Technical document analysis
-- Engineering insights
-- AI-powered problem solving
-- Cross-domain knowledge integration
-
-Ask me about engineering challenges, AI applications, or industry best practices!
-"""
-
-# UI Setup
-st.markdown("""
-    <style>
-    @import url('https://fonts.cdnfonts.com/css/tw-cen-mt');
-    * { font-family: 'Tw Cen MT', sans-serif; }
-    .st-emotion-cache-1y4p8pa { padding: 2rem 1rem; }
-    </style>
-""", unsafe_allow_html=True)
-
-st.title("🚀 Ataliba o Agent Nerdx 🚀")
-
-# Sidebar
+import tempfile
+from threading import Thread
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from langchain_community.document_loaders import PyPDFLoader, TextLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+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"]
+
+# --- Page Config ---
+st.set_page_config(page_title="Hybrid RAG with Streaming", page_icon="📄", layout="centered")
+st.title("📄 Hybrid Search + Streaming Chat")
+
+# --- Sidebar Upload ---
 with st.sidebar:
-    st.header("⚡️ Hugging Face Model Loaded")
-    st.markdown("Model: `amiguel/unsloth_finetune_test` with LoRA")
-    uploaded_file = st.file_uploader("Upload technical documents", type=["pdf", "docx", "xlsx", "xlsm"])
-
-# Session state
-if "file_context" not in st.session_state:
-    st.session_state.file_context = None
-if "chat_history" not in st.session_state:
-    st.session_state.chat_history = []
-
-# File parser
-def parse_file(file):
-    try:
-        if file.type == "application/pdf":
-            reader = PyPDF2.PdfReader(file)
-            return "\n".join([page.extract_text() for page in reader.pages])
-        elif file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
-            doc = Document(file)
-            return "\n".join([para.text for para in doc.paragraphs])
-        elif file.type in ["application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "application/vnd.ms-excel"]:
-            df = pd.read_excel(file)
-            return df.to_string()
-    except Exception as e:
-        st.error(f"Error processing file: {str(e)}")
-        return None
-
-# Process file
-if uploaded_file and not st.session_state.file_context:
-    st.session_state.file_context = parse_file(uploaded_file)
-    if st.session_state.file_context:
-        st.sidebar.success("✅ Document loaded successfully")
-
-# Load model
+    st.header("📤 Upload Documents")
+    uploaded_files = st.file_uploader("Upload PDFs or .txt files", type=["pdf", "txt"], accept_multiple_files=True)
+    clear_chat = st.button("🧹 Clear Conversation")
+
+# --- Session State ---
+if "messages" not in st.session_state or clear_chat:
+    st.session_state.messages = []
+
+# --- Load LLM ---
 @st.cache_resource
-def load_unsloth_model():
-    base_model = "unsloth/llama-3-8b-Instruct-bnb-4bit"
-    adapter = "amiguel/unsloth_finetune_test"
-    model, tokenizer = FastLanguageModel.from_pretrained(
-        model_name=base_model,
-        max_seq_length=2048,
-        dtype=None,
-        load_in_4bit=True
+def load_model():
+    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
+
+tokenizer, model = load_model()
+
+# --- Load & Chunk Documents ---
+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)
+        documents.extend(loader.load())
+    return documents
+
+def chunk_documents(documents):
+    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+    return splitter.split_documents(documents)
+
+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})
+
+    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 hybrid
+
+# --- 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"
     )
-    model.load_adapter(adapter)
-    FastLanguageModel.for_inference(model)
-    return model, tokenizer
-
-# Generate response
-def generate_response(prompt):
-    bio_triggers = ['who are you', 'ataliba', 'yourself', 'skilled at', 
-                    'background', 'experience', 'valonylabs', 'totalenergies']
-
-    if any(trigger in prompt.lower() for trigger in bio_triggers):
-        for line in ATALIBA_BIO.split('\n'):
-            yield line + '\n'
-            time.sleep(0.1)
-        return
-
-    try:
-        model, tokenizer = load_unsloth_model()
-        context = st.session_state.file_context or ""
-        full_prompt = f"You are an expert in life balance and general knowledge. Use the context to answer precisely.\nContext: {context}\n\nQuestion: {prompt}"
-
-        inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
-        outputs = model.generate(**inputs, max_new_tokens=256, do_sample=False)
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        for line in response.split('\n'):
-            yield line + '\n'
-            time.sleep(0.05)
-
-    except Exception as e:
-        yield f"⚠️ Model Error: {str(e)}"
-
-# Chat interface
-for msg in st.session_state.chat_history:
+    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):
+    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()
+    output = ""
+    for token in streamer:
+        output += token
+        yield output
+
+# --- Run Document Processing and Retrieval ---
+retriever = None
+if uploaded_files:
+    with st.spinner("📚 Processing documents..."):
+        docs = process_documents(uploaded_files)
+        chunks = chunk_documents(docs)
+        retriever = build_hybrid_retriever(chunks)
+    st.success("✅ Document processing complete.")
+
+# --- Display Past Messages ---
+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"])
 
-if prompt := st.chat_input("Ask about documents or technical matters..."):
-    st.session_state.chat_history.append({"role": "user", "content": prompt})
-    with st.chat_message("user", avatar=USER_AVATAR):
-        st.markdown(prompt)
+# --- 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})
 
-    with st.chat_message("assistant", avatar=BOT_AVATAR):
-        response_placeholder = st.empty()
-        full_response = ""
+    context = ""
+    if retriever:
+        docs = retriever.get_relevant_documents(prompt)
+        context = "\n\n".join([doc.page_content for doc in docs])
 
-        for chunk in generate_response(prompt):
-            full_response += chunk
-            response_placeholder.markdown(full_response + "▌")
+    full_prompt = build_prompt(st.session_state.messages, context=context)
 
-        response_placeholder.markdown(full_response)
-        st.session_state.chat_history.append({"role": "assistant", "content": full_response})
+    with st.chat_message("assistant", avatar=BOT_AVATAR):
+        response_container = st.empty()
+        answer = ""
+        for chunk in generate_streaming_response(full_prompt):
+            answer = chunk
+            response_container.markdown(answer + "▌", unsafe_allow_html=True)
+        response_container.markdown(answer)
+        st.session_state.messages.append({"role": "assistant", "content": answer})