Update app.py

app.py

@@ -1,3 +1,4 @@
+
 import streamlit as st
 import torch
 import os
@@ -8,7 +9,10 @@ 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
+from langchain.retrievers import EnsembleRetriever
 from langchain.schema import Document
+from langchain.docstore.document import Document as LangchainDocument
 
 # --- HF Token ---
 HF_TOKEN = st.secrets["HF_TOKEN"]
@@ -21,111 +25,83 @@ st.title("📂 DigiTs the Twin")
 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")
+    hybrid_toggle = st.checkbox("🔀 Enable Hybrid Search", value=True)
 
 # --- Model Loading ---
 @st.cache_resource
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("amiguel/GM_Qwen1.8B_Finetune", trust_remote_code=True, token=HF_TOKEN)
-    model = AutoModelForCausalLM.from_pretrained(
-        "amiguel/GM_Qwen1.8B_Finetune",
-        device_map="auto",
-        torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32,
-        trust_remote_code=True,
-        token=HF_TOKEN
-    )
-    return model, tokenizer
-
-model, tokenizer = load_model()
-
-# --- Prompt Helper ---
-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 "
-    "(with emphasis on Visual Internal Inspection - VII), and pressure safety devices (PSDs). Rely on uploaded documents "
-    "and context to provide practical, standards-driven, and technically accurate responses. Your guidance reflects deep "
-    "field experience, industry regulations, and proven methodologies in asset integrity and reliability engineering."
-)
+    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()
+
+# --- Document Processing ---
+def process_documents(files):
+    documents = []
+    for file in files:
+        if file.name.endswith(".pdf"):
+            loader = PyPDFLoader(file)
+        else:
+            loader = TextLoader(file)
+        docs = loader.load()
+        documents.extend(docs)
+    return documents
+
+def chunk_documents(documents):
+    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+    return splitter.split_documents(documents)
+
+# --- Embedding and Retrieval ---
+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
 
-def build_prompt(messages, context=""):
-    prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}\n\nContext:\n{context}<|im_end|>\n"
-    for msg in messages:
-        role = msg["role"]
-        prompt += f"<|im_start|>{role}\n{msg['content']}<|im_end|>\n"
-    prompt += "<|im_start|>assistant\n"
-    return prompt
+    ensemble = EnsembleRetriever(retrievers=[faiss_retriever, bm25_retriever], weights=[0.5, 0.5])
+    return faiss_retriever, ensemble
 
+# --- Inference ---
+def generate_answer(query, retriever):
+    docs = retriever.get_relevant_documents(query)
+    context = "\n".join([doc.page_content for doc in docs])
 
-# --- RAG Embedding and Search ---
-@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:
-            out_file.write(f.read())
-
-        loader = PyPDFLoader(file_path) if f.name.endswith(".pdf") else TextLoader(file_path)
-        raw_docs.extend(loader.load())
-
-    splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=64)
-    chunks = splitter.split_documents(raw_docs)
-    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-    db = FAISS.from_documents(chunks, embedding=embeddings)
-    return db
-
-retriever = embed_uploaded_files(uploaded_files) if uploaded_files else None
+    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). "
+        "Use the context below to answer professionally.\n\nContext:\n" + context + "\n\nQuery: " + query + "\nAnswer:"
+    )
 
-# --- Streaming Response ---
-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)
-    thread = Thread(target=model.generate, kwargs={
-        "input_ids": inputs["input_ids"],
-        "attention_mask": inputs["attention_mask"],
-        "max_new_tokens": 1024,
-        "temperature": 0.7,
-        "top_p": 0.9,
-        "repetition_penalty": 1.1,
-        "do_sample": True,
-        "streamer": streamer
-    })
+    inputs = tokenizer(system_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()
-    return streamer
-
-# --- Avatars & Messages ---
-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"
-
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-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):
-        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.similarity_search(prompt, k=3)
-        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):
-        start_time = time.time()
-        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})
+
+    answer = ""
+    for token in streamer:
+        answer += token
+        yield answer
+
+# --- Main App ---
+if uploaded_files:
+    with st.spinner("Processing documents..."):
+        docs = process_documents(uploaded_files)
+        chunks = chunk_documents(docs)
+        faiss_retriever, hybrid_retriever = build_retrievers(chunks)
+        st.success("Documents processed successfully.")
+
+    query = st.text_input("🔍 Ask a question based on the uploaded documents")
+    if query:
+        st.subheader("📤 Answer")
+        retriever = hybrid_retriever if hybrid_toggle else faiss_retriever
+        response_placeholder = st.empty()
+        full_response = ""
+        for partial_response in generate_answer(query, retriever):
+            full_response = partial_response
+            response_placeholder.markdown(full_response)