Update app.py

app.py

@@ -3,51 +3,35 @@ import gc
 import tempfile
 import uuid
 import pandas as pd
-
-from llama_index.llms.cerebras import Cerebras
+import streamlit as st
 from llama_index.core import Settings
+from llama_index.llms.cerebras import Cerebras
+from llama_index.core import PromptTemplate
 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
+from llama_index.core.llms import ChatMessage
 
 if "id" not in st.session_state:
     st.session_state.id = uuid.uuid4()
     st.session_state.file_cache = {}
 
 session_id = st.session_state.id
+client = None
 
-# Load Cerebras API key
-api_key = os.getenv("CEREBRAS_API_KEY")
-if not api_key:
-    raise ValueError("CEREBRAS_API_KEY is not set in Hugging Face Secrets.")
-else:
-    print("Cerebras API key loaded successfully.")
-
-@st.cache_resource
+# Initialize Cerebras LLM
 def load_llm():
-    # Initialize Cerebras client
-    os.environ["CEREBRAS_API_KEY"] = api_key
-    client = Cerebras(api_key=os.environ["CEREBRAS_API_KEY"])
-    Settings.client = client
-    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
+    # Replace with your API Key and model
+    api_key = os.getenv("CEREBRAS_API_KEY")
+    if not api_key:
+        api_key = st.text_input("Enter your Cerebras API key:")
+    llm = Cerebras(model="llama-3.3-70b", api_key=api_key)
+    return llm
 
 def reset_chat():
     st.session_state.messages = []
+    st.session_state.context = None
     gc.collect()
 
 def display_excel(file):
@@ -87,24 +71,42 @@ with st.sidebar:
 
                     docs = loader.load_data()
 
-                    # Setup embedding model
+                    # setup llm & embedding model
+                    llm = load_llm()  # Load the Cerebras 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
+                    Settings.embed_model = embed_model
                     node_parser = MarkdownNodeParser()
                     index = VectorStoreIndex.from_documents(documents=docs, transformations=[node_parser], show_progress=True)
 
-                    # Store index as query engine
+                    # Create the query engine
+                    Settings.llm = llm
                     query_engine = index.as_query_engine(streaming=True)
 
+                    # Customise prompt template
+                    qa_prompt_tmpl_str = (
+                        "Context information is below.\n"
+                        "---------------------\n"
+                        "{context_str}\n"
+                        "---------------------\n"
+                        "Given the context information above I want you to think step by step to answer the query in a highly precise and crisp manner focused on the final answer, in case you don't know the answer say 'I don't know!'.\n"
+                        "Query: {query_str}\n"
+                        "Answer: "
+                    )
+                    qa_prompt_tmpl = PromptTemplate(qa_prompt_tmpl_str)
+
+                    query_engine.update_prompts(
+                        {"response_synthesizer:text_qa_template": qa_prompt_tmpl}
+                    )
+
                     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
+                # Inform the user that the file is processed and Display the PDF uploaded
                 st.success("Ready to Chat!")
                 display_excel(uploaded_file)
+
         except Exception as e:
             st.error(f"An error occurred: {e}")
             st.stop()
@@ -112,7 +114,7 @@ with st.sidebar:
 col1, col2 = st.columns([6, 1])
 
 with col1:
-    st.header(f"RAG over Excel using DuckLink 🐥 & Llama-3.3-70B")
+    st.header(f"RAG over Excel using Dockling 🐥 &  Llama-3.3 70B")
 
 with col2:
     st.button("Clear ↺", on_click=reset_chat)
@@ -121,9 +123,6 @@ with col2:
 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"]):
@@ -142,13 +141,14 @@ if prompt := st.chat_input("What's up?"):
         message_placeholder = st.empty()
         full_response = ""
 
-        # Query Cerebras model
-        stream = query_cerebras(client, prompt)
+        # Using Cerebras stream_chat for streaming response
+        messages = [
+            ChatMessage(role="user", content=prompt)
+        ]
 
-        # Handle streaming response
-        for chunk in stream:
-            content = chunk.choices[0].delta.get("content", "")
-            full_response += content
+        response = llm.stream_chat(messages)
+        for r in response:
+            full_response += r.delta
             message_placeholder.markdown(full_response + "▌")
 
         message_placeholder.markdown(full_response)