Spaces:
Sleeping
Sleeping
| """ | |
| Qwen2.5 PDF RAG System for Hugging Face Spaces | |
| Adapted for deployment on Hugging Face Spaces with optimizations for the cloud environment. | |
| """ | |
| import os | |
| import time | |
| import gradio as gr | |
| from typing import List, Dict, Any, Tuple | |
| import torch | |
| # LangChain imports - updated to avoid deprecation warnings | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| from langchain_community.vectorstores import Chroma | |
| from langchain_community.embeddings import HuggingFaceEmbeddings | |
| from langchain_community.document_loaders import PyPDFLoader | |
| from langchain.schema import Document | |
| # Transformers for Qwen2.5 models (more compatible with HF Spaces) | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline | |
| import warnings | |
| warnings.filterwarnings("ignore") | |
| class PDFRagSystem: | |
| """PDF RAG System using Qwen2.5, ChromaDB, and LangChain - HF Spaces optimized""" | |
| def __init__(self, model_name: str = "Qwen/Qwen2.5-1.5B-Instruct", persist_directory: str = "db"): | |
| """ | |
| Initialize the RAG system | |
| Args: | |
| model_name: Name of the Qwen model to use | |
| persist_directory: Directory to store the vector database | |
| """ | |
| self.model_name = model_name | |
| self.persist_directory = persist_directory | |
| self.pipe = None | |
| self.tokenizer = None | |
| self.model = None | |
| self.vectorstore = None | |
| self.embeddings = None | |
| self.top_sources = [] | |
| # Check available device | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Using device: {self.device}") | |
| # Initialize embedding model | |
| print("Loading embedding model...") | |
| try: | |
| self.embeddings = HuggingFaceEmbeddings( | |
| model_name="sentence-transformers/all-MiniLM-L6-v2", | |
| model_kwargs={"device": self.device}, | |
| encode_kwargs={"normalize_embeddings": True} | |
| ) | |
| except Exception as e: | |
| print(f"Warning: Error loading HuggingFaceEmbeddings, trying alternative: {e}") | |
| # Fallback to basic embeddings if HuggingFaceEmbeddings fails | |
| from sentence_transformers import SentenceTransformer | |
| self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2') | |
| self.embeddings = self._create_custom_embeddings() | |
| # Load LLM | |
| self._load_llm() | |
| def _create_custom_embeddings(self): | |
| """Create custom embeddings wrapper if HuggingFaceEmbeddings fails""" | |
| class CustomEmbeddings: | |
| def __init__(self, model): | |
| self.model = model | |
| def embed_documents(self, texts): | |
| return self.model.encode(texts).tolist() | |
| def embed_query(self, text): | |
| return self.model.encode([text])[0].tolist() | |
| return CustomEmbeddings(self.embedding_model) | |
| def change_model(self, model_name: str) -> str: | |
| """ | |
| Change the LLM model | |
| Args: | |
| model_name: New model name to use | |
| Returns: | |
| Status message | |
| """ | |
| if self.model_name == model_name: | |
| return f"Already using model: {model_name}" | |
| self.model_name = model_name | |
| try: | |
| # Clear GPU memory | |
| if hasattr(self, 'model') and self.model is not None: | |
| del self.model | |
| del self.tokenizer | |
| del self.pipe | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| self._load_llm() | |
| return f"Successfully switched to model: {model_name}" | |
| except Exception as e: | |
| return f"Error switching model: {str(e)}" | |
| def _load_llm(self): | |
| """Load the Qwen2.5 model with optimized settings for HF Spaces""" | |
| print(f"\nLoading {self.model_name} model...") | |
| start_time = time.time() | |
| try: | |
| # Load tokenizer | |
| self.tokenizer = AutoTokenizer.from_pretrained( | |
| self.model_name, | |
| trust_remote_code=True | |
| ) | |
| # Configure model loading for limited resources | |
| model_kwargs = { | |
| "trust_remote_code": True, | |
| "torch_dtype": torch.float16 if self.device == "cuda" else torch.float32, | |
| "low_cpu_mem_usage": True, | |
| } | |
| if self.device == "cuda": | |
| model_kwargs["device_map"] = "auto" | |
| # Load model | |
| self.model = AutoModelForCausalLM.from_pretrained( | |
| self.model_name, | |
| **model_kwargs | |
| ) | |
| if self.device == "cpu": | |
| self.model = self.model.to(self.device) | |
| # Create pipeline | |
| self.pipe = pipeline( | |
| "text-generation", | |
| model=self.model, | |
| tokenizer=self.tokenizer, | |
| device=0 if self.device == "cuda" else -1, | |
| torch_dtype=torch.float16 if self.device == "cuda" else torch.float32, | |
| return_full_text=False | |
| ) | |
| load_time = time.time() - start_time | |
| print(f"Model loaded in {load_time:.2f} seconds") | |
| except Exception as e: | |
| print(f"Error loading model: {e}") | |
| # Fallback to a smaller model if the requested one fails | |
| if "1.5B" not in self.model_name: | |
| print("Falling back to Qwen2.5-1.5B-Instruct...") | |
| self.model_name = "Qwen/Qwen2.5-1.5B-Instruct" | |
| self._load_llm() | |
| else: | |
| raise e | |
| def process_pdf(self, pdf_file: str) -> List[Document]: | |
| """ | |
| Process a PDF file into documents for the vectorstore | |
| Args: | |
| pdf_file: Path to the PDF file | |
| Returns: | |
| List of document chunks | |
| """ | |
| try: | |
| loader = PyPDFLoader(pdf_file) | |
| documents = loader.load() | |
| # Split documents into chunks | |
| text_splitter = RecursiveCharacterTextSplitter( | |
| chunk_size=800, # Smaller chunks for better performance | |
| chunk_overlap=150, | |
| separators=["\n\n", "\n", ". ", " ", ""] | |
| ) | |
| chunks = text_splitter.split_documents(documents) | |
| return chunks | |
| except Exception as e: | |
| print(f"Error processing PDF {pdf_file}: {e}") | |
| return [] | |
| def create_vectorstore(self, pdf_files: List[str]) -> str: | |
| """ | |
| Create or update the vector store with documents from PDF files | |
| Args: | |
| pdf_files: List of paths to PDF files | |
| Returns: | |
| Status message | |
| """ | |
| if not pdf_files: | |
| return "No files provided." | |
| all_chunks = [] | |
| processed_files = 0 | |
| for pdf_file in pdf_files: | |
| if not os.path.exists(pdf_file): | |
| print(f"Warning: File {pdf_file} does not exist. Skipping.") | |
| continue | |
| print(f"Processing {pdf_file}...") | |
| chunks = self.process_pdf(pdf_file) | |
| if chunks: | |
| print(f"Created {len(chunks)} chunks from {pdf_file}") | |
| all_chunks.extend(chunks) | |
| processed_files += 1 | |
| else: | |
| print(f"Failed to process {pdf_file}") | |
| if not all_chunks: | |
| return "No documents were successfully processed." | |
| try: | |
| # Create or update vectorstore | |
| if os.path.exists(self.persist_directory) and len(os.listdir(self.persist_directory)) > 0: | |
| print("Loading existing vectorstore...") | |
| self.vectorstore = Chroma( | |
| persist_directory=self.persist_directory, | |
| embedding_function=self.embeddings | |
| ) | |
| print("Adding new documents to existing vectorstore...") | |
| self.vectorstore.add_documents(all_chunks) | |
| else: | |
| print("Creating new vectorstore...") | |
| self.vectorstore = Chroma.from_documents( | |
| documents=all_chunks, | |
| embedding=self.embeddings, | |
| persist_directory=self.persist_directory | |
| ) | |
| # Persist to disk | |
| self.vectorstore.persist() | |
| return f"Successfully processed {processed_files} PDFs with {len(all_chunks)} chunks." | |
| except Exception as e: | |
| return f"Error creating vectorstore: {str(e)}" | |
| def retrieve_context(self, query: str, k: int = 4) -> Tuple[str, List[Dict]]: | |
| """ | |
| Retrieve relevant context for a query | |
| Args: | |
| query: User query | |
| k: Number of top documents to retrieve | |
| Returns: | |
| Tuple of (concatenated context string, list of source documents) | |
| """ | |
| if not self.vectorstore: | |
| return "", [] | |
| try: | |
| # Search for relevant documents | |
| docs_with_scores = self.vectorstore.similarity_search_with_score(query, k=k) | |
| context_parts = [] | |
| sources = [] | |
| for i, (doc, score) in enumerate(docs_with_scores): | |
| context_part = f"Document {i+1}:\n{doc.page_content}\n" | |
| context_parts.append(context_part) | |
| # Clean metadata for serialization | |
| clean_metadata = {} | |
| for key, value in doc.metadata.items(): | |
| str_key = str(key) | |
| if isinstance(value, (str, int, float, bool, type(None))): | |
| clean_metadata[str_key] = value | |
| else: | |
| clean_metadata[str_key] = str(value) | |
| source_info = { | |
| "content": str(doc.page_content), | |
| "metadata": clean_metadata, | |
| "score": float(score), | |
| "source_id": i+1 | |
| } | |
| sources.append(source_info) | |
| self.top_sources = sources | |
| context = "\n".join(context_parts) | |
| return context, sources | |
| except Exception as e: | |
| print(f"Error retrieving context: {e}") | |
| return "", [] | |
| def generate_response(self, query: str, system_prompt: str = "You are a helpful assistant that answers questions based on the provided documents.") -> str: | |
| """ | |
| Generate a response using RAG | |
| Args: | |
| query: User query | |
| system_prompt: System prompt to set assistant behavior | |
| Returns: | |
| Model response | |
| """ | |
| # Retrieve relevant context | |
| context, _ = self.retrieve_context(query) | |
| if not context: | |
| return "No relevant documents found in the database. Please upload some PDF files first." | |
| # Create RAG prompt | |
| rag_prompt = f"""Based on the following context, please answer the question. If the answer is not in the context, say that you don't know. | |
| Context: | |
| {context} | |
| Question: {query} | |
| Answer:""" | |
| try: | |
| # Generate response | |
| print(f"Running inference for query: {query}") | |
| start_time = time.time() | |
| # Use the pipeline for generation | |
| response = self.pipe( | |
| rag_prompt, | |
| max_new_tokens=300, | |
| temperature=0.7, | |
| top_p=0.9, | |
| do_sample=True, | |
| pad_token_id=self.tokenizer.eos_token_id | |
| ) | |
| inference_time = time.time() - start_time | |
| print(f"Inference completed in {inference_time:.2f} seconds") | |
| # Extract the generated text | |
| if isinstance(response, list) and len(response) > 0: | |
| result = response[0].get('generated_text', '').strip() | |
| else: | |
| result = str(response).strip() | |
| return result if result else "I couldn't generate a response. Please try again." | |
| except Exception as e: | |
| print(f"Error generating response: {e}") | |
| return f"Error generating response: {str(e)}" | |
| def get_top_sources(self) -> List[Dict]: | |
| """Get the top sources used for the last query""" | |
| return self.top_sources | |
| class RagUI: | |
| """Gradio UI for the PDF RAG System - HF Spaces optimized""" | |
| def __init__(self, rag_system: PDFRagSystem): | |
| self.rag_system = rag_system | |
| self.interface = None | |
| # Define available models (optimized for HF Spaces) | |
| self.models = { | |
| "Qwen2.5-1.5B (Recommended)": "Qwen/Qwen2.5-1.5B-Instruct", | |
| "Qwen2.5-3B": "Qwen/Qwen2.5-3B-Instruct" | |
| } | |
| self.current_model = "Qwen2.5-1.5B (Recommended)" | |
| def _upload_files(self, files) -> str: | |
| """Handle file upload""" | |
| if not files: | |
| return "No files selected." | |
| try: | |
| file_paths = [f.name for f in files] | |
| return self.rag_system.create_vectorstore(file_paths) | |
| except Exception as e: | |
| return f"Error processing files: {str(e)}" | |
| def _switch_model(self, model_name: str) -> str: | |
| """Switch the model""" | |
| if model_name not in self.models: | |
| return f"Unknown model: {model_name}" | |
| full_model_name = self.models[model_name] | |
| self.current_model = model_name | |
| return self.rag_system.change_model(full_model_name) | |
| def _query(self, query: str, system_prompt: str) -> Tuple[str, str]: | |
| """Process a query""" | |
| if not query.strip(): | |
| return "Please enter a question.", "" | |
| response = self.rag_system.generate_response(query, system_prompt) | |
| sources = self.rag_system.get_top_sources() | |
| sources_html = self._format_source_display(sources) | |
| return response, sources_html | |
| def _format_source_display(self, sources: List[Dict]) -> str: | |
| """Format sources for display""" | |
| if not sources: | |
| return "<div class='source-container'>No sources available.</div>" | |
| html = "<div class='source-container'>" | |
| for i, source in enumerate(sources): | |
| try: | |
| if not isinstance(source, dict): | |
| continue | |
| metadata = source.get("metadata", {}) | |
| if not isinstance(metadata, dict): | |
| metadata = {} | |
| page_num = metadata.get("page", "Unknown") | |
| source_file = metadata.get("source", "Unknown") | |
| content = source.get("content", "No content available")[:500] + "..." # Limit content length | |
| score = source.get("score", 0.0) | |
| source_id = source.get("source_id", i+1) | |
| # Determine relevance class | |
| if score <= 0.5: # Lower is better for distance-based similarity | |
| relevance_class = "relevance-high" | |
| elif score <= 0.8: | |
| relevance_class = "relevance-medium" | |
| else: | |
| relevance_class = "relevance-low" | |
| html += f""" | |
| <div class="source-card"> | |
| <div class="source-header"> | |
| Source {source_id} (<span class="{relevance_class}">Distance: {score:.2f}</span>) | |
| </div> | |
| <div class="source-meta"> | |
| <strong>File:</strong> {os.path.basename(str(source_file))} | |
| </div> | |
| <div class="source-meta"> | |
| <strong>Page:</strong> {page_num} | |
| </div> | |
| <div class="source-content"> | |
| {content} | |
| </div> | |
| </div> | |
| """ | |
| except Exception as e: | |
| html += f'<div class="source-card">Error displaying source {i+1}: {str(e)}</div>' | |
| html += "</div>" | |
| return html | |
| def build_interface(self): | |
| """Build the Gradio interface""" | |
| # Custom CSS for better appearance | |
| css = """ | |
| .source-container { | |
| max-height: 600px; | |
| overflow-y: auto; | |
| padding: 10px; | |
| } | |
| .source-card { | |
| margin-bottom: 15px; | |
| padding: 12px; | |
| border: 1px solid #ddd; | |
| border-radius: 6px; | |
| background-color: #f8f9fa; | |
| box-shadow: 0 1px 3px rgba(0,0,0,0.1); | |
| } | |
| .source-header { | |
| font-size: 16px; | |
| font-weight: bold; | |
| margin-bottom: 8px; | |
| color: #333; | |
| } | |
| .source-meta { | |
| color: #666; | |
| margin-bottom: 6px; | |
| font-size: 14px; | |
| } | |
| .source-content { | |
| background-color: #fff; | |
| padding: 10px; | |
| border-radius: 4px; | |
| border-left: 3px solid #007bff; | |
| font-family: 'Segoe UI', sans-serif; | |
| line-height: 1.4; | |
| font-size: 14px; | |
| } | |
| .relevance-high { color: #28a745; font-weight: bold; } | |
| .relevance-medium { color: #ffc107; font-weight: bold; } | |
| .relevance-low { color: #dc3545; font-weight: bold; } | |
| """ | |
| with gr.Blocks(title="Qwen2.5 PDF RAG System", css=css) as interface: | |
| gr.Markdown(""" | |
| # 🤖 Qwen2.5 PDF RAG System | |
| Upload PDF documents and ask questions about their content using advanced AI. | |
| **⚡ Powered by Qwen2.5 Language Models** | |
| """) | |
| with gr.Tab("📚 Main Interface"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("### 🔧 Settings") | |
| # Model selection | |
| model_dropdown = gr.Dropdown( | |
| choices=list(self.models.keys()), | |
| value=self.current_model, | |
| label="AI Model", | |
| info="1.5B model recommended for stability" | |
| ) | |
| model_switch_btn = gr.Button("🔄 Switch Model", size="sm") | |
| model_status = gr.Textbox( | |
| label="Model Status", | |
| value=f"Using: {self.current_model}", | |
| interactive=False | |
| ) | |
| gr.Markdown("### 📄 Upload Documents") | |
| file_input = gr.File( | |
| file_count="multiple", | |
| file_types=[".pdf"], | |
| label="PDF Files" | |
| ) | |
| upload_button = gr.Button("📤 Process PDFs", variant="primary") | |
| upload_status = gr.Textbox( | |
| label="Status", | |
| interactive=False, | |
| placeholder="Upload status will appear here..." | |
| ) | |
| with gr.Column(scale=2): | |
| gr.Markdown("### 💬 Ask Questions") | |
| system_prompt = gr.Textbox( | |
| label="System Instructions", | |
| value="You are a helpful AI assistant. Answer questions based only on the provided documents. Be concise and cite relevant information.", | |
| lines=3 | |
| ) | |
| query_input = gr.Textbox( | |
| label="Your Question", | |
| placeholder="What would you like to know about your documents?", | |
| lines=2 | |
| ) | |
| query_button = gr.Button("🔍 Ask Question", variant="primary") | |
| answer_output = gr.Textbox( | |
| label="Answer", | |
| interactive=False, | |
| lines=8, | |
| placeholder="Answers will appear here..." | |
| ) | |
| with gr.Tab("📖 Sources"): | |
| gr.Markdown("### 📚 Reference Documents") | |
| gr.Markdown("View the source documents used to generate answers.") | |
| sources_display = gr.HTML( | |
| label="Sources", | |
| value="<p>No sources available yet. Ask a question first!</p>" | |
| ) | |
| with gr.Tab("ℹ️ Info"): | |
| gr.Markdown(""" | |
| ### About This System | |
| This is a **Retrieval-Augmented Generation (RAG)** system that: | |
| - 📤 **Processes PDF documents** and stores them in a vector database | |
| - 🔍 **Searches** for relevant content based on your questions | |
| - 🤖 **Generates answers** using Qwen2.5 language models | |
| - 📚 **Shows sources** so you can verify the information | |
| ### Available Models | |
| - **Qwen2.5-1.5B** ⚡ - Fast and efficient (Recommended for HF Spaces) | |
| - **Qwen2.5-3B** 🧠 - More capable but slower | |
| ### Tips for Best Results | |
| 1. 📄 Upload clear, text-based PDFs (not scanned images) | |
| 2. ❓ Ask specific questions rather than broad topics | |
| 3. 🔍 Check the "Sources" tab to see what documents were used | |
| 4. 🔄 Try rephrasing your question if you don't get good results | |
| ### Technical Details | |
| - **Vector Store**: ChromaDB with cosine similarity | |
| - **Embeddings**: sentence-transformers/all-MiniLM-L6-v2 | |
| - **Chunk Size**: 800 tokens with 150 token overlap | |
| - **Context Window**: Up to 4 most relevant document chunks | |
| """) | |
| # Event handlers | |
| upload_button.click( | |
| fn=self._upload_files, | |
| inputs=[file_input], | |
| outputs=[upload_status] | |
| ) | |
| query_button.click( | |
| fn=self._query, | |
| inputs=[query_input, system_prompt], | |
| outputs=[answer_output, sources_display] | |
| ) | |
| query_input.submit( | |
| fn=self._query, | |
| inputs=[query_input, system_prompt], | |
| outputs=[answer_output, sources_display] | |
| ) | |
| model_switch_btn.click( | |
| fn=self._switch_model, | |
| inputs=[model_dropdown], | |
| outputs=[model_status] | |
| ) | |
| self.interface = interface | |
| return interface | |
| def launch(self, **kwargs): | |
| """Launch the Gradio interface""" | |
| if not self.interface: | |
| self.build_interface() | |
| return self.interface.launch(**kwargs) | |
| # Initialize and launch the application | |
| def main(): | |
| """Main function optimized for Hugging Face Spaces""" | |
| print("🚀 Starting Qwen2.5 PDF RAG System...") | |
| print(f"📱 Device: {'GPU' if torch.cuda.is_available() else 'CPU'}") | |
| # Use the lightweight model by default for HF Spaces | |
| model_name = "Qwen/Qwen2.5-1.5B-Instruct" | |
| # Create RAG system | |
| try: | |
| rag_system = PDFRagSystem(model_name, persist_directory="chroma_db") | |
| # Create and launch UI | |
| ui = RagUI(rag_system) | |
| ui.launch( | |
| share=True, | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| show_error=True | |
| ) | |
| except Exception as e: | |
| print(f"❌ Error starting application: {e}") | |
| # Create a simple error interface | |
| def error_interface(): | |
| return "❌ Failed to initialize the RAG system. Please check the logs." | |
| error_app = gr.Interface( | |
| fn=error_interface, | |
| inputs=[], | |
| outputs="text", | |
| title="Error - Qwen2.5 PDF RAG System" | |
| ) | |
| error_app.launch() | |
| if __name__ == "__main__": | |
| main() |