Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,158 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+# Imports
 import streamlit as st
+import tempfile
+import os
 
-"""
-# Welcome to Streamlit!
+# LangChain imports
+from langchain_community.document_loaders import (
+    PyPDFLoader, 
+    TextLoader, 
+    UnstructuredWordDocumentLoader, 
+    CSVLoader
+)
+from langchain_community.vectorstores import FAISS
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_huggingface import HuggingFaceEndpoint
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnablePassthrough
+from langchain.embeddings.base import Embeddings
+
+# HuggingFace Inference client
+from huggingface_hub import InferenceClient
+
+# Load HF API key from Secrets
+HF_TOKEN = os.environ.get("HUGGINGFACE_API_KEY")
+if not HF_TOKEN:
+    st.error("HuggingFace API key not found. Add it as a Space Secret: 'HUGGINGFACE_API_KEY'")
+    st.stop()
+
+# Initialize HF Inference client
+hf_client = InferenceClient(provider="huggingface-inference", api_key=HF_TOKEN)
+
+# Embeddings wrapper using HF Inference API
+class HFInferenceEmbeddings(Embeddings):
+    def embed_documents(self, texts):
+        embeddings = []
+        for text in texts:
+            res = hf_client.feature_extraction(
+                model="intfloat/multilingual-e5-large-instruct",
+                inputs=text
+            )
+            embeddings.append(res[0])  # HF returns a list of vectors
+        return embeddings
+
+    def embed_query(self, text):
+        res = hf_client.feature_extraction(
+            model="intfloat/multilingual-e5-large-instruct",
+            inputs=text
+        )
+        return res[0]
+
+# Setup LLM via HF Endpoint
+llm = HuggingFaceEndpoint(
+    repo_id="AI-Sweden-Models/Llama-3-8B-instruct",
+    task="text-generation",
+    provider="huggingface-inference",
+    temperature=0.2,
+    max_new_tokens=512,
+    api_key=HF_TOKEN,
+)
+
+# Streamlit UI
+st.title("Ask RAG - Multi-file Support")
+
+uploaded_files = st.file_uploader(
+    "Upload files (PDF, DOCX, TXT, CSV)", 
+    type=["pdf", "docx", "txt", "csv"], 
+    accept_multiple_files=True
+)
+
+@st.cache_resource
+def load_files(files):
+    if not files:
+        return None, []
+
+    loaders = []
+    temp_files = []
+
+    for file in files:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.name)[-1]) as temp_file:
+            temp_file.write(file.read())
+            temp_path = temp_file.name
+            temp_files.append(temp_path)
 
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
+        if file.name.endswith(".pdf"):
+            loaders.append(PyPDFLoader(temp_path))
+        elif file.name.endswith(".txt"):
+            loaders.append(TextLoader(temp_path))
+        elif file.name.endswith(".docx"):
+            loaders.append(UnstructuredWordDocumentLoader(temp_path))
+        elif file.name.endswith(".csv"):
+            loaders.append(CSVLoader(temp_path))
 
-In the meantime, below is an example of what you can do with just a few lines of code:
+    documents = []
+    for loader in loaders:
+        documents.extend(loader.load())
+
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+    docs = text_splitter.split_documents(documents)
+
+    vectorstore = FAISS.from_documents(docs, HFInferenceEmbeddings())
+    return vectorstore, temp_files
+
+if uploaded_files:
+    vectorstore, temp_files = load_files(uploaded_files)
+else:
+    vectorstore, temp_files = None, []
+
+if vectorstore:
+    retriever = vectorstore.as_retriever()
+
+    chat_prompt = ChatPromptTemplate.from_template(
+        """Use the context below to answer the question.
+Context:
+{context}
+Question:
+{question}
 """
+    )
+
+    def format_docs(docs):
+        return "\n\n".join(d.page_content for d in docs)
+
+    rag_chain = (
+        {
+            "context": retriever | format_docs,
+            "question": RunnablePassthrough(),
+        }
+        | chat_prompt
+        | llm
+    )
+
+    if 'messages' not in st.session_state:
+        st.session_state.messages = []
+
+    for message in st.session_state.messages:
+        st.chat_message(message["role"]).markdown(message["content"])
+
+    user_input = st.chat_input("Enter your prompt")
+    if user_input:
+        st.chat_message("user").markdown(user_input)
+        st.session_state.messages.append({"role": "user", "content": user_input})
+
+        response = rag_chain.invoke(user_input)
+        answer = response.content
+        st.chat_message("assistant").markdown(answer)
+        st.session_state.messages.append({"role": "assistant", "content": answer})
+
+else:
+    st.warning("Please upload files to start querying.")
 
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+if st.button("Clear All"):
+    st.session_state.messages = []
+    for file_path in temp_files:
+        try:
+            os.remove(file_path)
+        except Exception as e:
+            print(f"Error deleting file {file_path}: {e}")
+    st.experimental_rerun()