Spaces:
Running
Running
| import os | |
| import tempfile | |
| import json | |
| import streamlit as st | |
| from streamlit_chat import message | |
| from preprocessing import Model | |
| from io import BytesIO | |
| import pickle | |
| # Home Page Setup | |
| st.set_page_config( | |
| page_title="PDF Insight Pro", | |
| page_icon="📄", | |
| layout="centered", | |
| ) | |
| # Custom CSS for a more polished look | |
| st.markdown(""" | |
| <style> | |
| .main { | |
| background-color: #f5f5f5; | |
| } | |
| .stButton button { | |
| background-color: #4CAF50; | |
| color: white; | |
| border-radius: 8px; | |
| } | |
| .stTextInput input { | |
| border-radius: 8px; | |
| padding: 10px; | |
| } | |
| .stFileUploader input { | |
| border-radius: 8px; | |
| } | |
| .stMarkdown h1 { | |
| color: #4CAF50; | |
| } | |
| </style> | |
| """, unsafe_allow_html=True) | |
| # Custom title and header | |
| st.title("📄 PDF Insight Pro") | |
| st.subheader("Empower Your Documents with AI-Driven Insights") | |
| def display_messages(): | |
| """ | |
| Displays the chat messages in the Streamlit UI. | |
| """ | |
| st.subheader("🗨️ Conversation") | |
| st.markdown("---") | |
| for i, (msg, is_user) in enumerate(st.session_state["messages"]): | |
| message(msg, is_user=is_user, key=str(i)) | |
| st.session_state["process_input_spinner"] = st.empty() | |
| def process_user_input(): | |
| """ | |
| Processes the user input by generating a response from the assistant. | |
| """ | |
| if st.session_state["user_input"] and len(st.session_state["user_input"].strip()) > 0: | |
| user_input = st.session_state["user_input"].strip() | |
| with st.session_state["process_input_spinner"], st.spinner("Analyzing..."): | |
| agent_response = st.session_state["assistant"].get_response( | |
| user_input, | |
| st.session_state["temperature"], | |
| st.session_state["max_tokens"], | |
| st.session_state["model"] | |
| ) | |
| st.session_state["messages"].append((user_input, True)) | |
| st.session_state["messages"].append((agent_response, False)) | |
| st.session_state["user_input"] = "" | |
| # Save chat history temporarily on local storage | |
| with open("chat_history.pkl", "wb") as f: | |
| pickle.dump(st.session_state["messages"], f) | |
| def process_file(): | |
| """ | |
| Processes the uploaded PDF file and appends its content to the context. | |
| """ | |
| for file in st.session_state["file_uploader"]: | |
| with tempfile.NamedTemporaryFile(delete=False) as tf: | |
| tf.write(file.getbuffer()) | |
| file_path = tf.name | |
| with st.session_state["process_file_spinner"], st.spinner(f"Processing {file.name}..."): | |
| try: | |
| st.session_state["assistant"].add_to_context(file_path) | |
| except Exception as e: | |
| st.error(f"Failed to process file {file.name}: {str(e)}") | |
| os.remove(file_path) | |
| def download_chat_history(): | |
| """ | |
| Allows users to download chat history in HTML or JSON format. | |
| """ | |
| # Convert messages to JSON format | |
| chat_data = [{"role": "user" if is_user else "assistant", "content": msg} for msg, is_user in st.session_state["messages"]] | |
| # Download as JSON | |
| json_data = json.dumps(chat_data, indent=4) | |
| st.download_button( | |
| label="💾 Download Chat History as JSON", | |
| data=json_data, | |
| file_name="chat_history.json", | |
| mime="application/json" | |
| ) | |
| # Download as HTML | |
| html_data = "<html><body><h1>Chat History</h1><ul>" | |
| for entry in chat_data: | |
| role = "User" if entry["role"] == "user" else "Assistant" | |
| html_data += f"<li><strong>{role}:</strong> {entry['content']}</li>" | |
| html_data += "</ul></body></html>" | |
| st.download_button( | |
| label="💾 Download Chat History as HTML", | |
| data=html_data, | |
| file_name="chat_history.html", | |
| mime="text/html" | |
| ) | |
| def main_page(): | |
| """ | |
| Main function to set up the Streamlit UI and handle user interactions. | |
| """ | |
| # Initialize session state variables | |
| if "messages" not in st.session_state: | |
| st.session_state["messages"] = [] | |
| if "assistant" not in st.session_state: | |
| st.session_state["assistant"] = Model() | |
| if "user_input" not in st.session_state: | |
| st.session_state["user_input"] = "" | |
| if "temperature" not in st.session_state: | |
| st.session_state["temperature"] = 0.5 | |
| if "max_tokens" not in st.session_state: | |
| st.session_state["max_tokens"] = 550 | |
| if "model" not in st.session_state: | |
| st.session_state["model"] = "llama-3.1-8b-instant" | |
| # File uploader | |
| st.subheader("📤 Upload Your PDF Documents") | |
| st.file_uploader( | |
| "Choose PDF files to analyze", | |
| type=["pdf"], | |
| key="file_uploader", | |
| on_change=process_file, | |
| accept_multiple_files=True, | |
| ) | |
| st.session_state["process_file_spinner"] = st.empty() | |
| # Document management section | |
| if st.session_state["assistant"].contexts: | |
| st.subheader("🗂️ Manage Uploaded Documents") | |
| for i, context in enumerate(st.session_state["assistant"].contexts): | |
| st.text_area(f"Document {i+1} Context", context[:500] + "..." if len(context) > 500 else context, height=100) | |
| if st.button(f"Remove Document {i+1}"): | |
| st.session_state["assistant"].remove_from_context(i) | |
| # Model settings | |
| with st.expander("⚙️ Customize AI Settings", expanded=True): | |
| st.slider("Sampling Temperature", min_value=0.0, max_value=1.0, step=0.1, key="temperature", help="Higher values make output more random.") | |
| st.slider("Max Tokens", min_value=750, max_value=5000, step=50, key="max_tokens", help="Limits the length of the response.") | |
| st.selectbox("Choose AI Model", ["llama-3.1-8b-instant", "llama3-70b-8192", "gemma-7b-it"], key="model") | |
| # Display messages and input box | |
| display_messages() | |
| st.text_input("Type your query and hit Enter", key="user_input", on_change=process_user_input, placeholder="Ask something about your documents...") | |
| # Download chat history section | |
| st.subheader("💾 Download Chat History") | |
| download_chat_history() | |
| # Developer info and bug report | |
| st.subheader("🐞 Bug Report") | |
| st.markdown(""" | |
| If you encounter any bugs or issues while using the app, please send a bug report to the developer. You can include a screenshot (optional) to help identify the problem.\n | |
| """) | |
| st.subheader("💡 Suggestions") | |
| st.markdown(""" | |
| Suggestions to improve the app's UI and user interface are also welcome. Feel free to reach out to the developer with your suggestions.\n | |
| """) | |
| st.subheader("👨💻 Developer Info") | |
| st.markdown(""" | |
| **Developer**: Jatin Mehra\n | |
| **Email**: [email protected]\n | |
| **Mobile**: 9910364780\n | |
| """) | |
| if __name__ == "__main__": | |
| main_page() |