Update app.py

app.py

@@ -1,102 +1,102 @@
-import os
-import streamlit as st
-# Update these imports
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain.chains import RetrievalQA
-from langchain_community.vectorstores import FAISS
-from langchain_core.prompts import PromptTemplate
-from langchain_huggingface import HuggingFaceEndpoint
-
-from dotenv import load_dotenv, find_dotenv
-load_dotenv(find_dotenv())
-
-DB_FAISS_PATH = "vectorstore/db_faiss"
-
-@st.cache_resource
-def get_vectorstore():
-    embedding_model = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
-    db = FAISS.load_local(DB_FAISS_PATH, embedding_model, allow_dangerous_deserialization=True)
-    return db
-
-def set_custom_prompt(custom_prompt_template):
-    prompt = PromptTemplate(template=custom_prompt_template, input_variables=["context", "question"])
-    return prompt
-
-def load_llm(huggingface_repo_id, HF_TOKEN):
-    llm = HuggingFaceEndpoint(
-        repo_id=huggingface_repo_id,
-        task="text-generation",  # Add this line
-        temperature=0.5,
-        model_kwargs={
-            "token": HF_TOKEN,
-            "max_length": 512  # Changed to integer
-        }
-    )
-    return llm
-
-def main():
-    st.title("Ask Chatbot!")
-    
-    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'])
-    
-    prompt = st.chat_input("Pass your prompt here")
-    
-    if prompt:
-        st.chat_message('user').markdown(prompt)
-        st.session_state.messages.append({'role': 'user', 'content': prompt})
-        
-        CUSTOM_PROMPT_TEMPLATE = """
-        Use the pieces of information provided in the context to answer user's question.
-        If you dont know the answer, just say that you dont know, dont try to make up an answer.
-        
-        Dont provide anything out of the given context
-        
-        Context: {context}
-        Question: {question}
-        
-        Start the answer directly. No small talk please.
-        """
-        
-        HUGGINGFACE_REPO_ID = "mistralai/Mistral-7B-Instruct-v0.3"
-        HF_TOKEN = os.environ.get("HF_TOKEN")
-        
-        try:
-            with st.spinner("Thinking..."):  # Add loading indicator
-                vectorstore = get_vectorstore()
-                if vectorstore is None:
-                    st.error("Failed to load the vector store")
-                    return
-
-                qa_chain = RetrievalQA.from_chain_type(
-                    llm=load_llm(huggingface_repo_id=HUGGINGFACE_REPO_ID, HF_TOKEN=HF_TOKEN),
-                    chain_type="stuff",
-                    retriever=vectorstore.as_retriever(search_kwargs={'k': 3}),
-                    return_source_documents=True,
-                    chain_type_kwargs={'prompt': set_custom_prompt(CUSTOM_PROMPT_TEMPLATE)}
-                )
-                
-                response = qa_chain.invoke({'query': prompt})
-                
-                result = response["result"]
-                source_documents = response["source_documents"]
-                
-                # Format source documents more cleanly
-                source_docs_text = "\n\n**Source Documents:**\n"
-                for i, doc in enumerate(source_documents, 1):
-                    source_docs_text += f"{i}. Page {doc.metadata.get('page', 'N/A')}: {doc.page_content[:200]}...\n\n"
-                
-                result_to_show = f"{result}\n{source_docs_text}"
-                
-                st.chat_message('assistant').markdown(result_to_show)
-                st.session_state.messages.append({'role': 'assistant', 'content': result_to_show})
-                
-        except Exception as e:
-            st.error(f"Error: {str(e)}")
-            st.error("Please check your HuggingFace token and model access permissions")
-
-if __name__ == "__main__":
+import os
+import streamlit as st
+# Update these imports
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.chains import RetrievalQA
+from langchain_community.vectorstores import FAISS
+from langchain_core.prompts import PromptTemplate
+from langchain_huggingface import HuggingFaceEndpoint
+
+from dotenv import load_dotenv, find_dotenv
+load_dotenv(find_dotenv())
+
+DB_FAISS_PATH = "vectorstore/db_faiss"
+
+@st.cache_resource
+def get_vectorstore():
+    embedding_model = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
+    db = FAISS.load_local(DB_FAISS_PATH, embedding_model, allow_dangerous_deserialization=True)
+    return db
+
+def set_custom_prompt(custom_prompt_template):
+    prompt = PromptTemplate(template=custom_prompt_template, input_variables=["context", "question"])
+    return prompt
+
+def load_llm(huggingface_repo_id, HF_TOKEN):
+    llm = HuggingFaceEndpoint(
+        repo_id=huggingface_repo_id,
+        task="text-generation",  # Add this line
+        temperature=0.5,
+        model_kwargs={
+            "token": HF_TOKEN,
+            "max_length": 512  # Changed to integer
+        }
+    )
+    return llm
+
+def main():
+    st.title("Forecasting discharge outcomes for critically ILL patients using machine learning!")
+    
+    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'])
+    
+    prompt = st.chat_input("Pass your prompt here")
+    
+    if prompt:
+        st.chat_message('user').markdown(prompt)
+        st.session_state.messages.append({'role': 'user', 'content': prompt})
+        
+        CUSTOM_PROMPT_TEMPLATE = """
+        Use the pieces of information provided in the context to answer user's question.
+        If you dont know the answer, just say that you dont know, dont try to make up an answer.
+        
+        Dont provide anything out of the given context
+        
+        Context: {context}
+        Question: {question}
+        
+        Start the answer directly. No small talk please.
+        """
+        
+        HUGGINGFACE_REPO_ID = "mistralai/Mistral-7B-Instruct-v0.3"
+        HF_TOKEN = os.environ.get("HF_TOKEN")
+        
+        try:
+            with st.spinner("Thinking..."):  # Add loading indicator
+                vectorstore = get_vectorstore()
+                if vectorstore is None:
+                    st.error("Failed to load the vector store")
+                    return
+
+                qa_chain = RetrievalQA.from_chain_type(
+                    llm=load_llm(huggingface_repo_id=HUGGINGFACE_REPO_ID, HF_TOKEN=HF_TOKEN),
+                    chain_type="stuff",
+                    retriever=vectorstore.as_retriever(search_kwargs={'k': 3}),
+                    return_source_documents=True,
+                    chain_type_kwargs={'prompt': set_custom_prompt(CUSTOM_PROMPT_TEMPLATE)}
+                )
+                
+                response = qa_chain.invoke({'query': prompt})
+                
+                result = response["result"]
+                source_documents = response["source_documents"]
+                
+                # Format source documents more cleanly
+                source_docs_text = "\n\n**Source Documents:**\n"
+                for i, doc in enumerate(source_documents, 1):
+                    source_docs_text += f"{i}. Page {doc.metadata.get('page', 'N/A')}: {doc.page_content[:200]}...\n\n"
+                
+                result_to_show = f"{result}\n{source_docs_text}"
+                
+                st.chat_message('assistant').markdown(result_to_show)
+                st.session_state.messages.append({'role': 'assistant', 'content': result_to_show})
+                
+        except Exception as e:
+            st.error(f"Error: {str(e)}")
+            st.error("Please check your HuggingFace token and model access permissions")
+
+if __name__ == "__main__":
     main()