File size: 11,691 Bytes
aad6058
ce97e36
9889ee0
 
 
 
 
 
 
 
 
 
 
ce97e36
aad6058
9889ee0
aad6058
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
# 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()