Create app.py

app.py

@@ -0,0 +1,148 @@
+import os
+import gc
+import tempfile
+import uuid
+import pandas as pd
+
+from cerebras.cloud.sdk import Cerebras
+from llama_index.core import Settings
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
+from llama_index.readers.docling import DoclingReader
+from llama_index.core.node_parser import MarkdownNodeParser
+
+import streamlit as st
+
+if "id" not in st.session_state:
+    st.session_state.id = uuid.uuid4()
+    st.session_state.file_cache = {}
+
+session_id = st.session_state.id
+
+@st.cache_resource
+def load_llm():
+    # Initialize Cerebras client
+    client = Cerebras(api_key=os.environ.get("CEREBRAS_API_KEY"))
+    return client
+
+def query_cerebras(client, prompt, max_tokens=1024, temperature=0.2, top_p=1):
+    # Query Cerebras model
+    stream = client.chat.completions.create(
+        messages=[{"role": "user", "content": prompt}],
+        model="llama-3.3-70b",
+        stream=True,
+        max_completion_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+    )
+    return stream
+
+def reset_chat():
+    st.session_state.messages = []
+    gc.collect()
+
+def display_excel(file):
+    st.markdown("### Excel Preview")
+    # Read the Excel file
+    df = pd.read_excel(file)
+    # Display the dataframe
+    st.dataframe(df)
+
+with st.sidebar:
+    st.header(f"Add your documents!")
+
+    uploaded_file = st.file_uploader("Choose your `.xlsx` file", type=["xlsx", "xls"])
+
+    if uploaded_file:
+        try:
+            with tempfile.TemporaryDirectory() as temp_dir:
+                file_path = os.path.join(temp_dir, uploaded_file.name)
+
+                with open(file_path, "wb") as f:
+                    f.write(uploaded_file.getvalue())
+
+                file_key = f"{session_id}-{uploaded_file.name}"
+                st.write("Indexing your document...")
+
+                if file_key not in st.session_state.get('file_cache', {}):
+
+                    if os.path.exists(temp_dir):
+                        reader = DoclingReader()
+                        loader = SimpleDirectoryReader(
+                            input_dir=temp_dir,
+                            file_extractor={".xlsx": reader},
+                        )
+                    else:
+                        st.error('Could not find the file you uploaded, please check again...')
+                        st.stop()
+
+                    docs = loader.load_data()
+
+                    # Setup embedding model
+                    embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-large-en-v1.5", trust_remote_code=True)
+                    Settings.embed_model = embed_model
+
+                    # Creating an index over loaded data
+                    node_parser = MarkdownNodeParser()
+                    index = VectorStoreIndex.from_documents(documents=docs, transformations=[node_parser], show_progress=True)
+
+                    # Store index as query engine
+                    query_engine = index.as_query_engine(streaming=True)
+
+                    st.session_state.file_cache[file_key] = query_engine
+                else:
+                    query_engine = st.session_state.file_cache[file_key]
+
+                # Inform the user that the file is processed and display the Excel uploaded
+                st.success("Ready to Chat!")
+                display_excel(uploaded_file)
+        except Exception as e:
+            st.error(f"An error occurred: {e}")
+            st.stop()
+
+col1, col2 = st.columns([6, 1])
+
+with col1:
+    st.header(f"RAG over Excel using Dockling 🐥 & Llama-3.3-70B")
+
+with col2:
+    st.button("Clear ↺", on_click=reset_chat)
+
+# Initialize chat history
+if "messages" not in st.session_state:
+    reset_chat()
+
+# Initialize LLM client
+client = load_llm()
+
+# Display chat messages from history on app rerun
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Accept user input
+if prompt := st.chat_input("What's up?"):
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    # Display user message in chat message container
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    # Display assistant response in chat message container
+    with st.chat_message("assistant"):
+        message_placeholder = st.empty()
+        full_response = ""
+
+        # Query Cerebras model
+        stream = query_cerebras(client, prompt)
+
+        # Handle streaming response
+        for chunk in stream:
+            content = chunk.choices[0].delta.get("content", "")
+            full_response += content
+            message_placeholder.markdown(full_response + "▌")
+
+        message_placeholder.markdown(full_response)
+
+    # Add assistant response to chat history
+    st.session_state.messages.append({"role": "assistant", "content": full_response})