Spaces:
Sleeping
Sleeping
| # Streamlit App for Multimodal RAG System - Deployment Optimized | |
| import os, pathlib | |
| cache_root = pathlib.Path("/tmp") | |
| (os.makedirs := None) # placeholder to satisfy linters if needed | |
| (pathlib.Path("/tmp/hf")).mkdir(parents=True, exist_ok=True) | |
| (pathlib.Path("/tmp/torch")).mkdir(parents=True, exist_ok=True) | |
| os.environ.setdefault("HF_HOME", "/tmp/hf") | |
| os.environ.setdefault("TRANSFORMERS_CACHE", "/tmp/hf") | |
| os.environ.setdefault("HUGGINGFACE_HUB_CACHE", "/tmp/hf") | |
| os.environ.setdefault("HF_DATASETS_CACHE", "/tmp/hf") | |
| os.environ.setdefault("SENTENCE_TRANSFORMERS_HOME", "/tmp/hf") | |
| os.environ.setdefault("XDG_CACHE_HOME", "/tmp") | |
| os.environ.setdefault("TORCH_HOME", "/tmp/torch") | |
| os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") | |
| import streamlit as st | |
| import os | |
| import tempfile | |
| from pathlib import Path | |
| from config import ( | |
| STREAMLIT_CONFIG, SUPPORTED_DOCUMENT_TYPES, | |
| SUPPORTED_IMAGE_TYPES, SUPPORTED_AUDIO_TYPES, | |
| MODELS_CONFIG, MAX_FILE_SIZE_MB, DEMO_ENABLED | |
| ) | |
| from multimodal_rag import MultimodalRAG | |
| from device_setup import get_gpu_memory_info, clear_gpu_memory | |
| def create_demo_content(): | |
| '''Create demo content for the system''' | |
| st.markdown("### 📋 Try These Examples:") | |
| col1, col2, col3 = st.columns(3) | |
| with col1: | |
| st.markdown("**📄 Documents**") | |
| st.markdown("- Upload PDF files") | |
| st.markdown("- Upload DOCX files") | |
| st.markdown("- Upload TXT files") | |
| with col2: | |
| st.markdown("**🖼️ Images**") | |
| st.markdown("- Upload PNG/JPG images") | |
| st.markdown("- Screenshots") | |
| st.markdown("- Diagrams & charts") | |
| with col3: | |
| st.markdown("**🎵 Audio**") | |
| st.markdown("- Upload MP3/WAV files") | |
| st.markdown("- Voice recordings") | |
| st.markdown("- Lecture audio") | |
| def display_system_info(rag_system): | |
| '''Display system information in sidebar''' | |
| with st.sidebar: | |
| st.markdown("### 🔧 System Status") | |
| # Device info | |
| if rag_system.has_gpu: | |
| st.success("🚀 GPU Acceleration Enabled") | |
| try: | |
| gpu_info = get_gpu_memory_info() | |
| if gpu_info: | |
| st.text(f"GPU Memory: {gpu_info}") | |
| except: | |
| pass | |
| else: | |
| st.info("💻 Running on CPU") | |
| # Model info | |
| st.markdown("### 🤖 Loaded Models") | |
| model_info = { | |
| "Text Encoder": MODELS_CONFIG["text_encoder"], | |
| "Vision Model": "CLIP ViT-Base", | |
| "Speech Model": "Whisper Base", | |
| "LLM": "GPT4All" if rag_system.llm else "Not Available" | |
| } | |
| for name, model in model_info.items(): | |
| status = "✅" if model != "Not Available" else "❌" | |
| st.text(f"{status} {name}: {model}") | |
| def main(): | |
| '''Main Streamlit interface''' | |
| # Configure page | |
| st.set_page_config( | |
| page_title=STREAMLIT_CONFIG["page_title"], | |
| layout=STREAMLIT_CONFIG["layout"], | |
| page_icon="🤖", | |
| initial_sidebar_state=STREAMLIT_CONFIG["initial_sidebar_state"] | |
| ) | |
| # Header | |
| st.title("🤖 Multimodal RAG System") | |
| st.markdown("### Upload documents, images, and audio files for AI-powered semantic search!") | |
| # Initialize RAG system | |
| if 'rag_system' not in st.session_state: | |
| with st.spinner("🔄 Loading AI models... This may take a few minutes on first run."): | |
| try: | |
| st.session_state.rag_system = MultimodalRAG() | |
| st.balloons() | |
| st.success("🎉 System loaded successfully!") | |
| except Exception as e: | |
| st.error(f"❌ Failed to load system: {e}") | |
| st.info("💡 Please refresh the page and try again") | |
| st.stop() | |
| rag = st.session_state.rag_system | |
| # Display system info in sidebar | |
| display_system_info(rag) | |
| # Main interface | |
| tab1, tab2, tab3 = st.tabs(["📤 Upload & Search", "📊 Knowledge Base", "ℹ️ About"]) | |
| with tab1: | |
| # File upload section | |
| st.markdown("### 📁 Upload Files") | |
| uploaded_files = st.file_uploader( | |
| "Choose files to add to knowledge base", | |
| type=SUPPORTED_DOCUMENT_TYPES + SUPPORTED_IMAGE_TYPES + SUPPORTED_AUDIO_TYPES, | |
| accept_multiple_files=True, | |
| help=f"Supported formats: {', '.join(SUPPORTED_DOCUMENT_TYPES + SUPPORTED_IMAGE_TYPES + SUPPORTED_AUDIO_TYPES)}" | |
| ) | |
| # Process uploaded files | |
| if uploaded_files: | |
| if st.button("🔄 Process Files", type="primary"): | |
| processed_count = 0 | |
| failed_count = 0 | |
| progress_bar = st.progress(0) | |
| status_text = st.empty() | |
| for i, uploaded_file in enumerate(uploaded_files): | |
| try: | |
| # Check file size | |
| file_size_mb = uploaded_file.size / (1024 * 1024) | |
| if file_size_mb > MAX_FILE_SIZE_MB: | |
| st.warning(f"⚠️ {uploaded_file.name} is too large ({file_size_mb:.1f}MB > {MAX_FILE_SIZE_MB}MB)") | |
| failed_count += 1 | |
| continue | |
| status_text.text(f"Processing {uploaded_file.name}...") | |
| # Save temporary file | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=Path(uploaded_file.name).suffix) as tmp_file: | |
| tmp_file.write(uploaded_file.getbuffer()) | |
| tmp_file_path = tmp_file.name | |
| # Process file | |
| success = rag.add_document(tmp_file_path, uploaded_file.name) | |
| if success: | |
| processed_count += 1 | |
| st.success(f"✅ Processed: {uploaded_file.name}") | |
| else: | |
| failed_count += 1 | |
| st.error(f"❌ Failed: {uploaded_file.name}") | |
| # Clean up | |
| os.unlink(tmp_file_path) | |
| except Exception as e: | |
| st.error(f"❌ Error processing {uploaded_file.name}: {e}") | |
| failed_count += 1 | |
| progress_bar.progress((i + 1) / len(uploaded_files)) | |
| status_text.text("Processing complete!") | |
| st.info(f"📊 Processed: {processed_count} files, Failed: {failed_count} files") | |
| # Search section | |
| st.markdown("### 🔍 Search Knowledge Base") | |
| query = st.text_input( | |
| "Enter your question or search query:", | |
| placeholder="e.g., What is the main topic of the uploaded document?", | |
| help="Ask questions about your uploaded content" | |
| ) | |
| col1, col2 = st.columns([3, 1]) | |
| with col1: | |
| top_k = st.slider("Number of results", 1, 10, 5) | |
| with col2: | |
| search_button = st.button("🔍 Search", type="primary") | |
| if search_button and query: | |
| with st.spinner("🔄 Searching..."): | |
| try: | |
| results = rag.search(query, top_k=top_k) | |
| if results: | |
| st.markdown("### 📋 Search Results") | |
| for i, result in enumerate(results, 1): | |
| with st.expander(f"Result {i}: {result['filename']} (Score: {result['score']:.3f})"): | |
| st.markdown(f"**Content:** {result['content']}") | |
| st.markdown(f"**Type:** {result['content_type']}") | |
| # Generate answer if LLM is available | |
| if rag.llm: | |
| st.markdown("### 🤖 AI Generated Answer") | |
| with st.spinner("🧠 Generating answer..."): | |
| try: | |
| answer = rag.generate_answer(query, results) | |
| st.markdown(answer) | |
| except Exception as e: | |
| st.error(f"❌ Error generating answer: {e}") | |
| else: | |
| st.info("💡 Upload a question to get AI-generated answers (LLM not available)") | |
| else: | |
| st.warning("🤔 No results found. Try uploading some files first or use different keywords.") | |
| except Exception as e: | |
| st.error(f"❌ Search failed: {e}") | |
| with tab2: | |
| # Knowledge base overview | |
| st.markdown("### 📚 Knowledge Base Overview") | |
| try: | |
| doc_count = len(rag.documents) | |
| st.metric("Total Documents", doc_count) | |
| if doc_count > 0: | |
| st.markdown("### 📄 Stored Documents") | |
| for i, doc in enumerate(rag.documents[:10]): # Show first 10 | |
| st.text(f"• {doc['filename']} ({doc['content_type']})") | |
| if doc_count > 10: | |
| st.text(f"... and {doc_count - 10} more documents") | |
| # Clear database option | |
| if st.button("🗑️ Clear Knowledge Base", help="Remove all documents"): | |
| if st.confirm("Are you sure you want to clear all documents?"): | |
| rag.clear_database() | |
| st.success("✅ Knowledge base cleared!") | |
| st.rerun() | |
| else: | |
| st.info("📭 No documents in knowledge base. Upload some files to get started!") | |
| if DEMO_ENABLED: | |
| create_demo_content() | |
| except Exception as e: | |
| st.error(f"❌ Error accessing knowledge base: {e}") | |
| with tab3: | |
| # About section | |
| st.markdown("### ℹ️ About This System") | |
| st.markdown(''' | |
| This **Multimodal RAG (Retrieval-Augmented Generation)** system allows you to: | |
| - 📄 **Process Documents**: Upload PDFs, DOCX, and text files | |
| - 🖼️ **Analyze Images**: Extract text and understand visual content | |
| - 🎵 **Transcribe Audio**: Convert speech to searchable text | |
| - 🔍 **Semantic Search**: Find relevant content using AI embeddings | |
| - 🤖 **AI Answers**: Get intelligent responses from a local LLM | |
| ### 🛠️ Technology Stack | |
| - **Frontend**: Streamlit | |
| - **Embeddings**: Sentence Transformers + CLIP | |
| - **Vector Search**: FAISS | |
| - **Speech Recognition**: OpenAI Whisper | |
| - **OCR**: EasyOCR | |
| - **Local LLM**: GPT4All | |
| ### 🚀 Features | |
| - ⚡ **GPU Acceleration** (when available) | |
| - 🔒 **Privacy First** (all processing runs locally/on your server) | |
| - 📱 **Responsive Design** | |
| - 🎯 **Easy to Use** | |
| ''') | |
| # System stats | |
| st.markdown("### 📊 System Information") | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| st.markdown("**Supported File Types:**") | |
| st.text(f"Documents: {', '.join(SUPPORTED_DOCUMENT_TYPES)}") | |
| st.text(f"Images: {', '.join(SUPPORTED_IMAGE_TYPES)}") | |
| st.text(f"Audio: {', '.join(SUPPORTED_AUDIO_TYPES)}") | |
| with col2: | |
| st.markdown("**Current Configuration:**") | |
| st.text(f"Max file size: {MAX_FILE_SIZE_MB}MB") | |
| st.text(f"Device: {'GPU' if rag.has_gpu else 'CPU'}") | |
| st.text(f"Models loaded: {len([m for m in [rag.text_encoder, rag.clip_model, rag.whisper_model, rag.ocr_reader, rag.llm] if m is not None])}/5") | |
| if __name__ == "__main__": | |
| main() | |