commit

src/streamlit_app.py

@@ -1,659 +1,497 @@
-# Enterprise AI Assistant with RAG using IBM Granite Models
-# Complete implementation with Streamlit interface
+# HF Spaces Optimized Enterprise AI Assistant with RAG
+# Handles permission issues and environment constraints
 
 import streamlit as st
 import os
-import pandas as pd
-import numpy as np
-from typing import List, Dict, Any, Optional
+import sys
 import tempfile
-import json
-from datetime import datetime
-import logging
+import shutil
 from pathlib import Path
+import logging
 
-# Core libraries for RAG
-import chromadb
-from sentence_transformers import SentenceTransformer
-import torch
-from transformers import (
-    AutoTokenizer, 
-    AutoModelForCausalLM, 
-    pipeline,
-    BitsAndBytesConfig
+# Configure logging for HF Spaces
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler(sys.stdout)]
 )
-
-# Document processing
-import PyPDF2
-from docx import Document
-import openpyxl
-from bs4 import BeautifulSoup
-import email
-from email.mime.text import MIMEText
-import chardet
-
-# Additional utilities
-import re
-from urllib.parse import urlparse
-import hashlib
-import pickle
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Configuration
-class Config:
+# Environment detection and configuration
+class HFSpacesConfig:
+    """Configuration optimized for Hugging Face Spaces"""
+    
+    # Detect if running on HF Spaces
+    IS_HF_SPACES = os.getenv("SPACE_ID") is not None
+    
+    # Model configurations (optimized for HF Spaces)
     GRANITE_MODEL_NAME = "ibm-granite/granite-3.1-8b-instruct"
-    GRANITE_GUARDIAN_MODEL = "ibm-granite/granite-guardian-3.2-5b"
     EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
-    CHUNK_SIZE = 1000
-    CHUNK_OVERLAP = 200
-    MAX_CONTEXT_LENGTH = 4000
-    TEMPERATURE = 0.7
-    MAX_NEW_TOKENS = 512
-    TOP_K = 5
-    SUPPORTED_FORMATS = ['.pdf', '.docx', '.xlsx', '.txt', '.csv', '.html', '.json', '.md']
-
-class DocumentProcessor:
-    """Handles document processing and text extraction"""
+    
+    # Reduced parameters for HF Spaces constraints
+    CHUNK_SIZE = 512
+    MAX_CONTEXT_LENGTH = 2048
+    MAX_NEW_TOKENS = 256
+    TOP_K = 3
     
     @staticmethod
-    def extract_text_from_pdf(file_path: str) -> str:
-        """Extract text from PDF files"""
+    def get_temp_dir():
+        """Get appropriate temporary directory"""
+        if HFSpacesConfig.IS_HF_SPACES:
+            # Use /tmp in HF Spaces
+            return "/tmp/chroma_db"
+        else:
+            return tempfile.mkdtemp(prefix="chroma_db_")
+    
+    @staticmethod
+    def setup_environment():
+        """Setup environment variables for HF Spaces"""
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+        os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
+        os.environ["TRANSFORMERS_VERBOSITY"] = "error"
+
+# Initialize environment
+HFSpacesConfig.setup_environment()
+
+# Import ML libraries after environment setup
+try:
+    import chromadb
+    from sentence_transformers import SentenceTransformer
+    import torch
+    from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+    import pandas as pd
+    import numpy as np
+    DEPENDENCIES_AVAILABLE = True
+except ImportError as e:
+    logger.error(f"Missing dependencies: {e}")
+    DEPENDENCIES_AVAILABLE = False
+
+# Lightweight document processors
+import PyPDF2
+from docx import Document
+import json
+import csv
+
+class SimpleVectorStore:
+    """Simplified vector store with better error handling"""
+    
+    def __init__(self):
+        self.documents = []
+        self.embeddings = []
+        self.metadata = []
+        self.embedding_model = None
+        self.chroma_client = None
+        self.chroma_collection = None
+        
+        self._initialize_storage()
+    
+    def _initialize_storage(self):
+        """Initialize storage with fallback options"""
+        # Try to initialize embedding model
         try:
-            with open(file_path, 'rb') as file:
-                reader = PyPDF2.PdfReader(file)
-                text = ""
-                for page in reader.pages:
-                    text += page.extract_text() + "\n"
-                return text
+            self.embedding_model = SentenceTransformer(HFSpacesConfig.EMBEDDING_MODEL_NAME)
+            logger.info("Embedding model loaded successfully")
         except Exception as e:
-            logger.error(f"Error extracting text from PDF: {e}")
-            return ""
+            logger.error(f"Failed to load embedding model: {e}")
+            return False
+        
+        # Try to initialize ChromaDB
+        try:
+            # Try in-memory first for HF Spaces
+            if HFSpacesConfig.IS_HF_SPACES:
+                self.chroma_client = chromadb.Client()
+                logger.info("Using in-memory ChromaDB for HF Spaces")
+            else:
+                # Try persistent storage locally
+                db_path = HFSpacesConfig.get_temp_dir()
+                Path(db_path).mkdir(parents=True, exist_ok=True)
+                self.chroma_client = chromadb.PersistentClient(path=db_path)
+                logger.info(f"Using persistent ChromaDB at: {db_path}")
+            
+            # Create collection
+            self.chroma_collection = self.chroma_client.create_collection(
+                name="enterprise_docs",
+                get_or_create=True
+            )
+            return True
+            
+        except Exception as e:
+            logger.warning(f"ChromaDB initialization failed: {e}, using simple storage")
+            return False
     
-    @staticmethod
-    def extract_text_from_docx(file_path: str) -> str:
-        """Extract text from DOCX files"""
+    def add_documents(self, texts, metadatas):
+        """Add documents to the vector store"""
+        if not self.embedding_model:
+            logger.error("Embedding model not available")
+            return False
+        
         try:
-            doc = Document(file_path)
-            text = ""
-            for paragraph in doc.paragraphs:
-                text += paragraph.text + "\n"
-            return text
+            # Generate embeddings
+            embeddings = self.embedding_model.encode(texts)
+            
+            if self.chroma_collection:
+                # Use ChromaDB
+                ids = [f"doc_{i}_{hash(text) % 10000}" for i, text in enumerate(texts)]
+                self.chroma_collection.add(
+                    documents=texts,
+                    embeddings=embeddings.tolist(),
+                    metadatas=metadatas,
+                    ids=ids
+                )
+            else:
+                # Use simple storage
+                self.documents.extend(texts)
+                self.embeddings.extend(embeddings)
+                self.metadata.extend(metadatas)
+            
+            logger.info(f"Added {len(texts)} documents to vector store")
+            return True
+            
         except Exception as e:
-            logger.error(f"Error extracting text from DOCX: {e}")
-            return ""
+            logger.error(f"Error adding documents: {e}")
+            return False
     
-    @staticmethod
-    def extract_text_from_xlsx(file_path: str) -> str:
-        """Extract text from Excel files"""
+    def search(self, query, k=3):
+        """Search for similar documents"""
+        if not self.embedding_model:
+            return []
+        
         try:
-            workbook = openpyxl.load_workbook(file_path)
-            text = ""
-            for sheet_name in workbook.sheetnames:
-                sheet = workbook[sheet_name]
-                text += f"Sheet: {sheet_name}\n"
-                for row in sheet.iter_rows(values_only=True):
-                    row_text = " | ".join([str(cell) if cell else "" for cell in row])
-                    if row_text.strip():
-                        text += row_text + "\n"
-                text += "\n"
-            return text
+            query_embedding = self.embedding_model.encode([query])
+            
+            if self.chroma_collection:
+                # Use ChromaDB search
+                results = self.chroma_collection.query(
+                    query_embeddings=query_embedding.tolist(),
+                    n_results=k
+                )
+                
+                search_results = []
+                for i in range(len(results['documents'][0])):
+                    search_results.append({
+                        'text': results['documents'][0][i],
+                        'metadata': results['metadatas'][0][i],
+                        'score': 1 - results['distances'][0][i] if results['distances'] else 0.5
+                    })
+                return search_results
+            
+            else:
+                # Use simple cosine similarity
+                if not self.embeddings:
+                    return []
+                
+                from sklearn.metrics.pairwise import cosine_similarity
+                similarities = cosine_similarity(query_embedding, self.embeddings)[0]
+                
+                # Get top k results
+                top_indices = similarities.argsort()[-k:][::-1]
+                
+                results = []
+                for idx in top_indices:
+                    results.append({
+                        'text': self.documents[idx],
+                        'metadata': self.metadata[idx],
+                        'score': similarities[idx]
+                    })
+                return results
+                
         except Exception as e:
-            logger.error(f"Error extracting text from XLSX: {e}")
-            return ""
+            logger.error(f"Search error: {e}")
+            return []
+    
+    def get_stats(self):
+        """Get storage statistics"""
+        if self.chroma_collection:
+            try:
+                count = self.chroma_collection.count()
+                return {'count': count, 'type': 'ChromaDB'}
+            except:
+                return {'count': 0, 'type': 'ChromaDB (Error)'}
+        else:
+            return {'count': len(self.documents), 'type': 'Simple Storage'}
+
+class SimpleDocumentProcessor:
+    """Simplified document processor"""
     
     @staticmethod
-    def extract_text_from_csv(file_path: str) -> str:
-        """Extract text from CSV files"""
+    def process_pdf(file):
+        """Process PDF file"""
         try:
-            df = pd.read_csv(file_path)
-            return df.to_string()
+            reader = PyPDF2.PdfReader(file)
+            text = ""
+            for page in reader.pages:
+                text += page.extract_text() + "\n"
+            return text
         except Exception as e:
-            logger.error(f"Error extracting text from CSV: {e}")
+            logger.error(f"PDF processing error: {e}")
             return ""
     
     @staticmethod
-    def extract_text_from_html(file_path: str) -> str:
-        """Extract text from HTML files"""
+    def process_docx(file):
+        """Process DOCX file"""
         try:
-            with open(file_path, 'r', encoding='utf-8') as file:
-                soup = BeautifulSoup(file.read(), 'html.parser')
-                return soup.get_text()
+            doc = Document(file)
+            text = "\n".join([para.text for para in doc.paragraphs])
+            return text
         except Exception as e:
-            logger.error(f"Error extracting text from HTML: {e}")
+            logger.error(f"DOCX processing error: {e}")
             return ""
     
     @staticmethod
-    def extract_text_from_txt(file_path: str) -> str:
-        """Extract text from TXT files"""
+    def process_txt(file):
+        """Process text file"""
         try:
-            # Detect encoding
-            with open(file_path, 'rb') as file:
-                raw_data = file.read()
-                encoding = chardet.detect(raw_data)['encoding']
-            
-            with open(file_path, 'r', encoding=encoding) as file:
-                return file.read()
+            return file.read().decode('utf-8')
         except Exception as e:
-            logger.error(f"Error extracting text from TXT: {e}")
+            logger.error(f"TXT processing error: {e}")
             return ""
     
     @staticmethod
-    def extract_text_from_json(file_path: str) -> str:
-        """Extract text from JSON files"""
+    def process_csv(file):
+        """Process CSV file"""
         try:
-            with open(file_path, 'r') as file:
-                data = json.load(file)
-                return json.dumps(data, indent=2)
+            content = file.read().decode('utf-8')
+            return content
         except Exception as e:
-            logger.error(f"Error extracting text from JSON: {e}")
+            logger.error(f"CSV processing error: {e}")
             return ""
-    
-    def process_document(self, file_path: str) -> str:
-        """Process document based on file extension"""
-        file_extension = Path(file_path).suffix.lower()
-        
-        extractors = {
-            '.pdf': self.extract_text_from_pdf,
-            '.docx': self.extract_text_from_docx,
-            '.xlsx': self.extract_text_from_xlsx,
-            '.csv': self.extract_text_from_csv,
-            '.html': self.extract_text_from_html,
-            '.txt': self.extract_text_from_txt,
-            '.md': self.extract_text_from_txt,
-            '.json': self.extract_text_from_json,
-        }
-        
-        extractor = extractors.get(file_extension)
-        if extractor:
-            return extractor(file_path)
-        else:
-            logger.warning(f"Unsupported file format: {file_extension}")
-            return ""
-
-class TextChunker:
-    """Handles text chunking for RAG"""
-    
-    def __init__(self, chunk_size: int = Config.CHUNK_SIZE, chunk_overlap: int = Config.CHUNK_OVERLAP):
-        self.chunk_size = chunk_size
-        self.chunk_overlap = chunk_overlap
-    
-    def chunk_text(self, text: str, document_name: str = "") -> List[Dict[str, Any]]:
-        """Split text into chunks with metadata"""
-        chunks = []
-        sentences = re.split(r'[.!?]+', text)
-        
-        current_chunk = ""
-        current_length = 0
-        
-        for sentence in sentences:
-            sentence = sentence.strip()
-            if not sentence:
-                continue
-                
-            sentence_length = len(sentence)
-            
-            if current_length + sentence_length > self.chunk_size and current_chunk:
-                # Save current chunk
-                chunks.append({
-                    'text': current_chunk.strip(),
-                    'metadata': {
-                        'document_name': document_name,
-                        'chunk_id': len(chunks),
-                        'timestamp': datetime.now().isoformat()
-                    }
-                })
-                
-                # Start new chunk with overlap
-                overlap_text = current_chunk[-self.chunk_overlap:] if len(current_chunk) > self.chunk_overlap else current_chunk
-                current_chunk = overlap_text + " " + sentence
-                current_length = len(current_chunk)
-            else:
-                current_chunk += " " + sentence
-                current_length += sentence_length
-        
-        # Add the last chunk
-        if current_chunk.strip():
-            chunks.append({
-                'text': current_chunk.strip(),
-                'metadata': {
-                    'document_name': document_name,
-                    'chunk_id': len(chunks),
-                    'timestamp': datetime.now().isoformat()
-                }
-            })
-        
-        return chunks
 
-class VectorStore:
-    """Handles vector storage and retrieval using ChromaDB"""
-    
-    def __init__(self, collection_name: str = "enterprise_documents"):
-        self.client = chromadb.PersistentClient(path="./chroma_db")
-        self.collection_name = collection_name
-        self.embedding_model = SentenceTransformer(Config.EMBEDDING_MODEL_NAME)
-        
-        # Create or get collection
-        try:
-            self.collection = self.client.get_collection(collection_name)
-        except:
-            self.collection = self.client.create_collection(
-                name=collection_name,
-                metadata={"description": "Enterprise document embeddings"}
-            )
-    
-    def add_documents(self, chunks: List[Dict[str, Any]]) -> None:
-        """Add document chunks to vector store"""
-        texts = [chunk['text'] for chunk in chunks]
-        metadatas = [chunk['metadata'] for chunk in chunks]
-        
-        # Generate embeddings
-        embeddings = self.embedding_model.encode(texts).tolist()
-        
-        # Generate IDs
-        ids = [f"doc_{i}_{hashlib.md5(text.encode()).hexdigest()[:8]}" 
-               for i, text in enumerate(texts)]
-        
-        # Add to collection
-        self.collection.add(
-            documents=texts,
-            embeddings=embeddings,
-            metadatas=metadatas,
-            ids=ids
-        )
-        
-        logger.info(f"Added {len(chunks)} chunks to vector store")
-    
-    def similarity_search(self, query: str, k: int = Config.TOP_K) -> List[Dict[str, Any]]:
-        """Search for similar documents"""
-        query_embedding = self.embedding_model.encode([query]).tolist()
-        
-        results = self.collection.query(
-            query_embeddings=query_embedding,
-            n_results=k,
-            include=['documents', 'metadatas', 'distances']
-        )
-        
-        search_results = []
-        for i in range(len(results['documents'][0])):
-            search_results.append({
-                'text': results['documents'][0][i],
-                'metadata': results['metadatas'][0][i],
-                'distance': results['distances'][0][i] if results['distances'] else 0
-            })
-        
-        return search_results
-    
-    def get_collection_stats(self) -> Dict[str, Any]:
-        """Get statistics about the collection"""
-        count = self.collection.count()
-        return {
-            'total_documents': count,
-            'collection_name': self.collection_name
-        }
-
-class GraniteModel:
-    """Handles IBM Granite model loading and inference"""
+class SimpleRAGAssistant:
+    """Simplified RAG Assistant for HF Spaces"""
     
     def __init__(self):
+        self.vector_store = SimpleVectorStore()
+        self.doc_processor = SimpleDocumentProcessor()
         self.model = None
         self.tokenizer = None
-        self.guardian_pipeline = None
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    def load_model(self):
+        """Load model with error handling"""
+        if self.model is not None:
+            return True
         
-    @st.cache_resource
-    def load_model(_self):
-        """Load Granite model with caching"""
         try:
-            # Configure for efficient loading
-            if _self.device == "cuda":
-                quantization_config = BitsAndBytesConfig(
-                    load_in_4bit=True,
-                    bnb_4bit_compute_dtype=torch.float16,
-                    bnb_4bit_use_double_quant=True,
-                    bnb_4bit_quant_type="nf4"
-                )
-            else:
-                quantization_config = None
+            # Try to load a smaller model first
+            model_name = "microsoft/DialoGPT-medium"  # Fallback model
             
-            # Load tokenizer
-            _self.tokenizer = AutoTokenizer.from_pretrained(
-                Config.GRANITE_MODEL_NAME,
-                trust_remote_code=True
-            )
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+            self.model = AutoModelForCausalLM.from_pretrained(model_name)
             
-            # Load model
-            _self.model = AutoModelForCausalLM.from_pretrained(
-                Config.GRANITE_MODEL_NAME,
-                quantization_config=quantization_config,
-                device_map="auto" if _self.device == "cuda" else None,
-                torch_dtype=torch.float16 if _self.device == "cuda" else torch.float32,
-                trust_remote_code=True
-            )
-            
-            # Load Guardian model for safety
-            try:
-                _self.guardian_pipeline = pipeline(
-                    "text-classification",
-                    model=Config.GRANITE_GUARDIAN_MODEL,
-                    device=0 if _self.device == "cuda" else -1
-                )
-                logger.info("Granite Guardian model loaded successfully")
-            except Exception as e:
-                logger.warning(f"Could not load Guardian model: {e}")
-                _self.guardian_pipeline = None
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
             
-            logger.info(f"Granite model loaded successfully on {_self.device}")
+            logger.info(f"Loaded model: {model_name}")
             return True
             
         except Exception as e:
-            logger.error(f"Error loading Granite model: {e}")
+            logger.error(f"Model loading failed: {e}")
             return False
     
-    def check_safety(self, text: str) -> bool:
-        """Check if text is safe using Guardian model"""
-        if not self.guardian_pipeline:
-            return True  # If no guardian model, assume safe
+    def process_documents(self, files):
+        """Process uploaded documents"""
+        results = {'success': [], 'errors': []}
         
-        try:
-            result = self.guardian_pipeline(text)
-            # Assuming Guardian returns safety classification
-            return result[0]['label'].lower() == 'safe'
-        except Exception as e:
-            logger.warning(f"Error in safety check: {e}")
-            return True  # Default to safe if error
+        for file in files:
+            try:
+                file_extension = Path(file.name).suffix.lower()
+                
+                if file_extension == '.pdf':
+                    text = self.doc_processor.process_pdf(file)
+                elif file_extension == '.docx':
+                    text = self.doc_processor.process_docx(file)
+                elif file_extension in ['.txt', '.md']:
+                    text = self.doc_processor.process_txt(file)
+                elif file_extension == '.csv':
+                    text = self.doc_processor.process_csv(file)
+                else:
+                    results['errors'].append(f"Unsupported format: {file.name}")
+                    continue
+                
+                if text.strip():
+                    # Simple chunking
+                    chunks = self._chunk_text(text, file.name)
+                    texts = [chunk['text'] for chunk in chunks]
+                    metadatas = [chunk['metadata'] for chunk in chunks]
+                    
+                    if self.vector_store.add_documents(texts, metadatas):
+                        results['success'].append({
+                            'name': file.name,
+                            'chunks': len(chunks)
+                        })
+                    else:
+                        results['errors'].append(f"Failed to store: {file.name}")
+                else:
+                    results['errors'].append(f"No text extracted: {file.name}")
+                    
+            except Exception as e:
+                results['errors'].append(f"Error processing {file.name}: {str(e)}")
+        
+        return results
     
-    def generate_response(self, prompt: str, context: str = "") -> str:
-        """Generate response using Granite model"""
-        if not self.model or not self.tokenizer:
-            if not self.load_model():
-                return "Error: Could not load the model. Please check your setup."
+    def _chunk_text(self, text, filename):
+        """Simple text chunking"""
+        chunk_size = HFSpacesConfig.CHUNK_SIZE
+        chunks = []
         
-        # Safety check
-        if not self.check_safety(prompt):
-            return "I cannot provide a response to that query due to safety concerns."
+        sentences = text.split('.')
+        current_chunk = ""
         
-        # Construct the full prompt
-        system_prompt = """You are an Enterprise AI Assistant with access to company documents and policies. 
-        Provide helpful, accurate, and professional responses based on the provided context.
-        If you cannot answer based on the context, say so clearly."""
+        for sentence in sentences:
+            if len(current_chunk + sentence) < chunk_size:
+                current_chunk += sentence + "."
+            else:
+                if current_chunk:
+                    chunks.append({
+                        'text': current_chunk.strip(),
+                        'metadata': {'source': filename, 'chunk_id': len(chunks)}
+                    })
+                current_chunk = sentence + "."
         
-        if context:
-            full_prompt = f"{system_prompt}\n\nContext:\n{context}\n\nUser Question: {prompt}\n\nAssistant:"
-        else:
-            full_prompt = f"{system_prompt}\n\nUser Question: {prompt}\n\nAssistant:"
+        if current_chunk:
+            chunks.append({
+                'text': current_chunk.strip(),
+                'metadata': {'source': filename, 'chunk_id': len(chunks)}
+            })
         
-        try:
-            # Tokenize input
-            inputs = self.tokenizer.encode(full_prompt, return_tensors='pt')
-            
-            # Truncate if too long
-            if inputs.shape[1] > Config.MAX_CONTEXT_LENGTH:
-                inputs = inputs[:, -Config.MAX_CONTEXT_LENGTH:]
-            
-            inputs = inputs.to(self.device)
-            
-            # Generate response
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    inputs,
-                    max_new_tokens=Config.MAX_NEW_TOKENS,
-                    temperature=Config.TEMPERATURE,
-                    do_sample=True,
-                    pad_token_id=self.tokenizer.eos_token_id,
-                    eos_token_id=self.tokenizer.eos_token_id,
-                    repetition_penalty=1.1
-                )
-            
-            # Decode response
-            response = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
-            return response.strip()
-            
-        except Exception as e:
-            logger.error(f"Error generating response: {e}")
-            return f"I apologize, but I encountered an error while generating a response: {str(e)}"
-
-class EnterpriseRAGAssistant:
-    """Main RAG Assistant class"""
+        return chunks
     
-    def __init__(self):
-        self.doc_processor = DocumentProcessor()
-        self.text_chunker = TextChunker()
-        self.vector_store = VectorStore()
-        self.granite_model = GraniteModel()
+    def answer_query(self, query):
+        """Answer query using RAG"""
+        # Search for relevant documents
+        search_results = self.vector_store.search(query, k=HFSpacesConfig.TOP_K)
         
-    def process_and_store_documents(self, uploaded_files) -> Dict[str, Any]:
-        """Process uploaded files and store in vector database"""
-        results = {
-            'processed_files': [],
-            'errors': [],
-            'total_chunks': 0
-        }
+        if not search_results:
+            return {
+                'response': "I don't have enough information to answer your question. Please upload some documents first.",
+                'sources': []
+            }
+        
+        # Prepare context
+        context = "\n\n".join([result['text'][:200] + "..." for result in search_results])
         
-        for uploaded_file in uploaded_files:
+        # Generate response (simplified)
+        if self.model and self.tokenizer:
             try:
-                # Save uploaded file temporarily
-                with tempfile.NamedTemporaryFile(delete=False, suffix=Path(uploaded_file.name).suffix) as tmp_file:
-                    tmp_file.write(uploaded_file.read())
-                    tmp_file_path = tmp_file.name
+                prompt = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:"
+                inputs = self.tokenizer.encode(prompt, return_tensors='pt', max_length=512, truncation=True)
                 
-                # Extract text
-                text = self.doc_processor.process_document(tmp_file_path)
-                
-                if text:
-                    # Chunk text
-                    chunks = self.text_chunker.chunk_text(text, uploaded_file.name)
-                    
-                    # Store in vector database
-                    self.vector_store.add_documents(chunks)
-                    
-                    results['processed_files'].append({
-                        'name': uploaded_file.name,
-                        'chunks': len(chunks),
-                        'text_length': len(text)
-                    })
-                    results['total_chunks'] += len(chunks)
-                else:
-                    results['errors'].append(f"Could not extract text from {uploaded_file.name}")
+                with torch.no_grad():
+                    outputs = self.model.generate(
+                        inputs, 
+                        max_length=inputs.shape[1] + 100,
+                        num_return_sequences=1,
+                        temperature=0.7,
+                        pad_token_id=self.tokenizer.eos_token_id
+                    )
                 
-                # Clean up temporary file
-                os.unlink(tmp_file_path)
+                response = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
                 
             except Exception as e:
-                results['errors'].append(f"Error processing {uploaded_file.name}: {str(e)}")
-        
-        return results
-    
-    def answer_query(self, query: str) -> Dict[str, Any]:
-        """Answer user query using RAG"""
-        # Retrieve relevant documents
-        search_results = self.vector_store.similarity_search(query)
-        
-        # Prepare context
-        context = "\n\n".join([result['text'] for result in search_results])
-        
-        # Generate response
-        response = self.granite_model.generate_response(query, context)
+                logger.error(f"Generation error: {e}")
+                response = f"Based on the available documents: {context[:300]}..."
+        else:
+            # Fallback response
+            response = f"Based on the available documents, here's what I found: {context[:300]}..."
         
         return {
             'response': response,
-            'sources': search_results,
-            'context_used': bool(context)
+            'sources': search_results
         }
 
 def main():
     """Main Streamlit application"""
     st.set_page_config(
-        page_title="Enterprise AI Assistant with RAG",
+        page_title="Enterprise RAG Assistant (HF Spaces)",
         page_icon="🏢",
-        layout="wide",
-        initial_sidebar_state="expanded"
+        layout="wide"
     )
     
-    # Custom CSS
-    st.markdown("""
-    <style>
-    .main-header {
-        font-size: 2.5rem;
-        color: #1f4e79;
-        text-align: center;
-        margin-bottom: 2rem;
-    }
-    .stButton > button {
-        background-color: #0f62fe;
-        color: white;
-        font-weight: bold;
-    }
-    .success-box {
-        padding: 1rem;
-        border-radius: 0.5rem;
-        background-color: #d4edda;
-        border: 1px solid #c3e6cb;
-        color: #155724;
-    }
-    .error-box {
-        padding: 1rem;
-        border-radius: 0.5rem;
-        background-color: #f8d7da;
-        border: 1px solid #f5c6cb;
-        color: #721c24;
-    }
-    </style>
-    """, unsafe_allow_html=True)
+    # Check dependencies
+    if not DEPENDENCIES_AVAILABLE:
+        st.error("Some dependencies are missing. Please check the requirements.txt file.")
+        st.stop()
+    
+    st.title("🏢 Enterprise RAG Assistant")
+    st.caption("Optimized for Hugging Face Spaces")
     
     # Initialize session state
     if 'rag_assistant' not in st.session_state:
-        st.session_state.rag_assistant = EnterpriseRAGAssistant()
+        with st.spinner("Initializing assistant..."):
+            st.session_state.rag_assistant = SimpleRAGAssistant()
     
     if 'chat_history' not in st.session_state:
         st.session_state.chat_history = []
     
-    # Header
-    st.markdown('<h1 class="main-header">🏢 Enterprise AI Assistant with RAG</h1>', unsafe_allow_html=True)
-    st.markdown("**Powered by IBM Granite Models | Intelligent Document Processing & Q&A**")
-    
     # Sidebar
     with st.sidebar:
-        st.header("📁 Document Management")
+        st.header("📁 Document Upload")
         
-        # File upload
         uploaded_files = st.file_uploader(
-            "Upload Enterprise Documents",
-            type=['pdf', 'docx', 'xlsx', 'txt', 'csv', 'html', 'json', 'md'],
-            accept_multiple_files=True,
-            help="Upload documents to build your knowledge base"
+            "Upload documents",
+            type=['pdf', 'docx', 'txt', 'csv', 'md'],
+            accept_multiple_files=True
         )
         
-        if uploaded_files and st.button("Process Documents", type="primary"):
+        if uploaded_files and st.button("Process Documents"):
             with st.spinner("Processing documents..."):
-                results = st.session_state.rag_assistant.process_and_store_documents(uploaded_files)
+                results = st.session_state.rag_assistant.process_documents(uploaded_files)
             
-            if results['processed_files']:
-                st.markdown('<div class="success-box">', unsafe_allow_html=True)
-                st.success(f"Successfully processed {len(results['processed_files'])} files!")
-                st.write(f"Total chunks created: {results['total_chunks']}")
-                
-                for file_info in results['processed_files']:
-                    st.write(f"✓ {file_info['name']}: {file_info['chunks']} chunks")
-                st.markdown('</div>', unsafe_allow_html=True)
+            if results['success']:
+                st.success(f"✅ Processed {len(results['success'])} files")
+                for file in results['success']:
+                    st.write(f"- {file['name']}: {file['chunks']} chunks")
             
             if results['errors']:
-                st.markdown('<div class="error-box">', unsafe_allow_html=True)
-                st.error("Some files had errors:")
+                st.error("❌ Some files had errors:")
                 for error in results['errors']:
-                    st.write(f"✗ {error}")
-                st.markdown('</div>', unsafe_allow_html=True)
-        
-        # Database stats
-        st.header("📊 Knowledge Base Stats")
-        try:
-            stats = st.session_state.rag_assistant.vector_store.get_collection_stats()
-            st.metric("Total Documents", stats['total_documents'])
-        except:
-            st.metric("Total Documents", 0)
-        
-        # Model info
-        st.header("🤖 Model Information")
-        st.info(f"**Main Model**: {Config.GRANITE_MODEL_NAME}")
-        st.info(f"**Safety Model**: {Config.GRANITE_GUARDIAN_MODEL}")
-        st.info(f"**Embedding Model**: {Config.EMBEDDING_MODEL_NAME}")
+                    st.write(f"- {error}")
+        
+        # Stats
+        st.header("📊 Statistics")
+        stats = st.session_state.rag_assistant.vector_store.get_stats()
+        st.metric("Documents", stats['count'])
+        st.info(f"Storage: {stats['type']}")
+        
+        # Model loading
+        st.header("🤖 Model Status")
+        if st.button("Load Model"):
+            with st.spinner("Loading model..."):
+                success = st.session_state.rag_assistant.load_model()
+            if success:
+                st.success("✅ Model loaded")
+            else:
+                st.error("❌ Model loading failed")
     
-    # Main content area
-    col1, col2 = st.columns([2, 1])
+    # Main chat interface
+    st.header("💬 Chat")
     
-    with col1:
-        st.header("💬 Chat with Your Documents")
-        
-        # Chat interface
-        query = st.text_input(
-            "Ask a question about your documents:",
-            placeholder="e.g., What is our company's policy on remote work?",
-            key="user_query"
-        )
-        
-        if st.button("Send Query", type="primary") and query:
-            with st.spinner("Generating response..."):
-                result = st.session_state.rag_assistant.answer_query(query)
-            
-            # Add to chat history
-            st.session_state.chat_history.append({
-                'query': query,
-                'response': result['response'],
-                'sources': result['sources'],
-                'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-            })
-        
-        # Display chat history
-        if st.session_state.chat_history:
-            st.header("📜 Chat History")
-            
-            for i, chat in enumerate(reversed(st.session_state.chat_history)):
-                with st.expander(f"Q: {chat['query'][:50]}... ({chat['timestamp']})", expanded=i==0):
-                    st.markdown("**Question:**")
-                    st.write(chat['query'])
-                    
-                    st.markdown("**Answer:**")
-                    st.write(chat['response'])
-                    
-                    if chat['sources']:
-                        st.markdown("**Sources:**")
-                        for j, source in enumerate(chat['sources'][:3]):
-                            st.markdown(f"**Source {j+1}** (from {source['metadata']['document_name']}):")
-                            st.text(source['text'][:200] + "...")
+    query = st.text_input("Ask a question about your documents:")
     
-    with col2:
-        st.header("🔍 Search Results")
-        
-        if st.session_state.chat_history:
-            latest_chat = st.session_state.chat_history[-1]
+    if st.button("Send") and query:
+        with st.spinner("Generating response..."):
+            result = st.session_state.rag_assistant.answer_query(query)
+        
+        st.session_state.chat_history.append({
+            'query': query,
+            'response': result['response'],
+            'sources': result['sources']
+        })
+    
+    # Display chat history
+    for i, chat in enumerate(reversed(st.session_state.chat_history)):
+        with st.expander(f"Q: {chat['query'][:50]}...", expanded=i==0):
+            st.write("**Question:**", chat['query'])
+            st.write("**Answer:**", chat['response'])
             
-            st.subheader("Latest Query Sources")
-            for i, source in enumerate(latest_chat['sources']):
-                with st.expander(f"Source {i+1}: {source['metadata']['document_name']}"):
-                    st.write(f"**Relevance Score**: {1 - source['distance']:.3f}")
-                    st.write(f"**Document**: {source['metadata']['document_name']}")
-                    st.write(f"**Chunk ID**: {source['metadata']['chunk_id']}")
-                    st.text_area("Content", source['text'], height=150, disabled=True)
-        
-        # Quick actions
-        st.header("⚡ Quick Actions")
-        if st.button("Clear Chat History"):
-            st.session_state.chat_history = []
-            st.rerun()
-        
-        if st.button("Reset Knowledge Base"):
-            if st.confirm("Are you sure you want to reset the knowledge base? This cannot be undone."):
-                # This would require implementing a reset method
-                st.warning("Knowledge base reset functionality would be implemented here")
+            if chat['sources']:
+                st.write("**Sources:**")
+                for j, source in enumerate(chat['sources'][:2]):
+                    st.write(f"{j+1}. {source['metadata']['source']} (Score: {source['score']:.2f})")
+                    st.text(source['text'][:150] + "...")
     
     # Footer
     st.markdown("---")
-    st.markdown(
-        "Built with ❤️ using IBM Granite Models, Streamlit, and ChromaDB | "
-        "Enterprise-grade AI Assistant for document processing and intelligent Q&A"
-    )
+    st.markdown("🤗 Running on Hugging Face Spaces | Built with Streamlit")
 
 if __name__ == "__main__":
     main()