import gradio as gr import torch from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline from sentence_transformers import SentenceTransformer import faiss import numpy as np import PyPDF2 import docx import io import os import re from typing import List, Optional, Dict, Tuple import json from collections import Counter import warnings warnings.filterwarnings("ignore") class SmartDocumentRAG: def __init__(self): print("🚀 Initializing Enhanced Smart RAG System...") # Initialize better embedding model self.embedder = SentenceTransformer('all-MiniLM-L6-v2') # Faster and good quality print("✅ Embedding model loaded") # Initialize optimized LLM with better quantization self.setup_llm() # Document storage self.documents = [] self.document_metadata = [] self.index = None self.is_indexed = False self.raw_text = "" self.document_type = "general" self.document_summary = "" self.sentence_embeddings = [] self.sentences = [] def setup_llm(self): """Setup optimized model with better quantization""" try: # Check CUDA availability device = "cuda" if torch.cuda.is_available() else "cpu" print(f"🔧 Using device: {device}") if device == "cuda": self.setup_gpu_model() else: self.setup_cpu_model() except Exception as e: print(f"❌ Error loading models: {e}") self.setup_fallback_model() def setup_gpu_model(self): """Setup GPU model with proper quantization""" try: # Use Phi-2 - excellent for Q&A and reasoning model_name = "microsoft/DialoGPT-medium" # Better quantization config quantization_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_quant_storage=torch.uint8 ) try: # Try Flan-T5 first - excellent for Q&A model_name = "google/flan-t5-base" print(f"🤖 Loading {model_name}...") self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModelForCausalLM.from_pretrained( model_name, quantization_config=quantization_config, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True ) # Create pipeline for easier use self.qa_pipeline = pipeline( "text2text-generation", model=self.model, tokenizer=self.tokenizer, max_length=512, do_sample=True, temperature=0.3, top_p=0.9 ) print("✅ Flan-T5 model loaded successfully") self.model_type = "flan-t5" except Exception as e: print(f"Flan-T5 failed, trying Phi-2: {e}") # Try Phi-2 as backup model_name = "microsoft/phi-2" print(f"🤖 Loading {model_name}...") self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) self.model = AutoModelForCausalLM.from_pretrained( model_name, quantization_config=quantization_config, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True ) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token print("✅ Phi-2 model loaded successfully") self.model_type = "phi-2" except Exception as e: print(f"❌ GPU models failed: {e}") self.setup_cpu_model() def setup_cpu_model(self): """Setup CPU-optimized model""" try: # Use DistilBERT for Q&A - much better than DialoGPT for this task model_name = "distilbert-base-cased-distilled-squad" print(f"🤖 Loading CPU model: {model_name}") self.qa_pipeline = pipeline( "question-answering", model=model_name, tokenizer=model_name ) self.model_type = "distilbert-qa" print("✅ DistilBERT Q&A model loaded successfully") except Exception as e: print(f"❌ CPU model failed: {e}") self.setup_fallback_model() def setup_fallback_model(self): """Fallback to basic model""" try: print("🤖 Loading fallback model...") self.qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad") self.model_type = "fallback" print("✅ Fallback model loaded") except Exception as e: print(f"❌ All models failed: {e}") self.qa_pipeline = None self.model_type = "none" def detect_document_type(self, text: str) -> str: """Enhanced document type detection""" text_lower = text.lower() resume_patterns = [ 'experience', 'skills', 'education', 'linkedin', 'email', 'phone', 'work experience', 'employment', 'resume', 'cv', 'curriculum vitae', 'internship', 'projects', 'achievements', 'career', 'profile', 'objective' ] research_patterns = [ 'abstract', 'introduction', 'methodology', 'conclusion', 'references', 'literature review', 'hypothesis', 'study', 'research', 'findings', 'data analysis', 'results', 'discussion', 'bibliography', 'journal' ] business_patterns = [ 'company', 'revenue', 'market', 'strategy', 'business', 'financial', 'quarter', 'profit', 'sales', 'growth', 'investment', 'stakeholder', 'operations', 'management', 'corporate', 'enterprise', 'budget' ] technical_patterns = [ 'implementation', 'algorithm', 'system', 'technical', 'specification', 'architecture', 'development', 'software', 'programming', 'api', 'database', 'framework', 'deployment', 'infrastructure', 'code' ] def count_matches(patterns, text): score = 0 for pattern in patterns: count = text.count(pattern) score += count * (2 if len(pattern.split()) > 1 else 1) # Weight phrases higher return score scores = { 'resume': count_matches(resume_patterns, text_lower), 'research': count_matches(research_patterns, text_lower), 'business': count_matches(business_patterns, text_lower), 'technical': count_matches(technical_patterns, text_lower) } max_score = max(scores.values()) if max_score > 5: # Higher threshold return max(scores, key=scores.get) return 'general' def create_document_summary(self, text: str) -> str: """Enhanced document summary creation""" try: clean_text = re.sub(r'\s+', ' ', text).strip() sentences = re.split(r'[.!?]+', clean_text) sentences = [s.strip() for s in sentences if len(s.strip()) > 30] if not sentences: return "Document contains basic information." # Use first few sentences and key information if self.document_type == 'resume': return self.extract_resume_summary(sentences, clean_text) elif self.document_type == 'research': return self.extract_research_summary(sentences) elif self.document_type == 'business': return self.extract_business_summary(sentences) else: return self.extract_general_summary(sentences) except Exception as e: print(f"Summary creation error: {e}") return "Document summary not available." def extract_resume_summary(self, sentences: List[str], full_text: str) -> str: """Extract resume-specific summary with better name detection""" summary_parts = [] # Extract name using multiple patterns name = self.extract_name(full_text) if name: summary_parts.append(f"Resume of {name}") # Extract role/title role_patterns = [ r'(?:software|senior|junior|lead|principal)?\s*(?:engineer|developer|analyst|manager|designer|architect|consultant)', r'(?:full stack|frontend|backend|data|ml|ai)\s*(?:engineer|developer)', r'(?:product|project|technical)\s*manager' ] for sentence in sentences[:5]: for pattern in role_patterns: matches = re.findall(pattern, sentence.lower()) if matches: summary_parts.append(f"working as {matches[0].title()}") break # Extract experience exp_match = re.search(r'(\d+)[\+\-\s]*(?:years?|yrs?)\s*(?:of\s*)?(?:experience|exp)', full_text.lower()) if exp_match: summary_parts.append(f"with {exp_match.group(1)}+ years of experience") return '. '.join(summary_parts) + '.' if summary_parts else "Professional resume with career details." def extract_name(self, text: str) -> str: """Extract name from document using multiple strategies""" # Strategy 1: Look for name patterns at the beginning lines = text.split('\n')[:10] # First 10 lines for line in lines: line = line.strip() if len(line) < 50 and len(line) > 3: # Likely a header line # Check if it looks like a name name_match = re.match(r'^([A-Z][a-z]+\s+[A-Z][a-z]+(?:\s+[A-Z][a-z]+)?)(?:\s|$)', line) if name_match: return name_match.group(1) # Strategy 2: Look for "Name:" pattern name_patterns = [ r'(?:name|full name):\s*([A-Z][a-z]+\s+[A-Z][a-z]+(?:\s+[A-Z][a-z]+)?)', r'^([A-Z][a-z]+\s+[A-Z][a-z]+)(?:\s*\n|\s*email|\s*phone|\s*linkedin)', ] for pattern in name_patterns: match = re.search(pattern, text, re.MULTILINE | re.IGNORECASE) if match: return match.group(1) return "" def extract_research_summary(self, sentences: List[str]) -> str: """Extract research paper summary""" # Look for abstract or introduction for sentence in sentences[:5]: if any(word in sentence.lower() for word in ['abstract', 'study', 'research', 'paper']): return sentence[:200] + ('...' if len(sentence) > 200 else '') return "Research document with academic content." def extract_business_summary(self, sentences: List[str]) -> str: """Extract business document summary""" for sentence in sentences[:3]: if any(word in sentence.lower() for word in ['company', 'business', 'organization']): return sentence[:200] + ('...' if len(sentence) > 200 else '') return "Business document with organizational information." def extract_general_summary(self, sentences: List[str]) -> str: """Extract general document summary""" return sentences[0][:200] + ('...' if len(sentences[0]) > 200 else '') if sentences else "General document." def extract_text_from_file(self, file_path: str) -> str: """Enhanced text extraction""" try: file_extension = os.path.splitext(file_path)[1].lower() if file_extension == '.pdf': return self.extract_from_pdf(file_path) elif file_extension == '.docx': return self.extract_from_docx(file_path) elif file_extension == '.txt': return self.extract_from_txt(file_path) else: return f"Unsupported file format: {file_extension}" except Exception as e: return f"Error reading file: {str(e)}" def extract_from_pdf(self, file_path: str) -> str: """Enhanced PDF extraction""" text = "" try: with open(file_path, 'rb') as file: pdf_reader = PyPDF2.PdfReader(file) for page in pdf_reader.pages: page_text = page.extract_text() if page_text.strip(): # Better text cleaning page_text = re.sub(r'\s+', ' ', page_text) page_text = re.sub(r'([a-z])([A-Z])', r'\1 \2', page_text) # Fix merged words text += f"{page_text}\n" except Exception as e: text = f"Error reading PDF: {str(e)}" return text.strip() def extract_from_docx(self, file_path: str) -> str: """Enhanced DOCX extraction""" try: doc = docx.Document(file_path) text = "" for paragraph in doc.paragraphs: if paragraph.text.strip(): text += paragraph.text.strip() + "\n" return text.strip() except Exception as e: return f"Error reading DOCX: {str(e)}" def extract_from_txt(self, file_path: str) -> str: """Enhanced TXT extraction""" encodings = ['utf-8', 'latin-1', 'cp1252', 'iso-8859-1'] for encoding in encodings: try: with open(file_path, 'r', encoding=encoding) as file: return file.read().strip() except UnicodeDecodeError: continue except Exception as e: return f"Error reading TXT: {str(e)}" return "Error: Could not decode file" def enhanced_chunk_text(self, text: str) -> List[Dict]: """Enhanced chunking with better overlap""" if not text.strip(): return [] chunks = [] # Split into sentences sentences = re.split(r'[.!?]+', text) sentences = [s.strip() for s in sentences if len(s.strip()) > 20] self.sentences = sentences # Create overlapping chunks chunk_size = 4 # sentences per chunk overlap = 2 # sentence overlap for i in range(0, len(sentences), chunk_size - overlap): chunk_sentences = sentences[i:i + chunk_size] if chunk_sentences: chunk_text = '. '.join(chunk_sentences) + '.' chunks.append({ 'text': chunk_text, 'sentence_indices': list(range(i, min(i + chunk_size, len(sentences)))), 'doc_type': self.document_type }) return chunks def process_documents(self, files) -> str: """Enhanced document processing""" if not files: return "❌ No files uploaded!" try: all_text = "" processed_files = [] for file in files: if file is None: continue file_text = self.extract_text_from_file(file.name) if not file_text.startswith("Error") and not file_text.startswith("Unsupported"): all_text += f"\n{file_text}" processed_files.append(os.path.basename(file.name)) else: return f"❌ {file_text}" if not all_text.strip(): return "❌ No text extracted from files!" # Store and analyze self.raw_text = all_text self.document_type = self.detect_document_type(all_text) self.document_summary = self.create_document_summary(all_text) # Enhanced chunking chunk_data = self.enhanced_chunk_text(all_text) if not chunk_data: return "❌ No valid text chunks created!" self.documents = [chunk['text'] for chunk in chunk_data] self.document_metadata = chunk_data # Create embeddings print(f"📄 Creating embeddings for {len(self.documents)} chunks...") embeddings = self.embedder.encode(self.documents, show_progress_bar=False) # Build FAISS index dimension = embeddings.shape[1] self.index = faiss.IndexFlatIP(dimension) # Normalize for cosine similarity faiss.normalize_L2(embeddings) self.index.add(embeddings.astype('float32')) self.is_indexed = True return f"✅ Successfully processed {len(processed_files)} files:\n" + \ f"📄 Files: {', '.join(processed_files)}\n" + \ f"📊 Document Type: {self.document_type.title()}\n" + \ f"🔍 Created {len(self.documents)} chunks\n" + \ f"📝 Summary: {self.document_summary}\n" + \ f"🚀 Ready for Q&A!" except Exception as e: return f"❌ Error processing documents: {str(e)}" def find_relevant_content(self, query: str, k: int = 3) -> str: """Improved content retrieval""" if not self.is_indexed: return "" try: # Semantic search query_embedding = self.embedder.encode([query]) faiss.normalize_L2(query_embedding) scores, indices = self.index.search(query_embedding.astype('float32'), min(k, len(self.documents))) relevant_chunks = [] for i, idx in enumerate(indices[0]): if idx < len(self.documents) and scores[0][i] > 0.1: # Lower threshold relevant_chunks.append(self.documents[idx]) return ' '.join(relevant_chunks) except Exception as e: print(f"Error in content retrieval: {e}") return "" def answer_question(self, query: str) -> str: """Enhanced question answering with better model usage""" if not query.strip(): return "❓ Please ask a question!" if not self.is_indexed: return "📁 Please upload and process documents first!" try: query_lower = query.lower() # Handle summary requests if any(word in query_lower for word in ['summary', 'summarize', 'about', 'overview']): return f"📄 **Document Summary:**\n\n{self.document_summary}" # Get relevant content context = self.find_relevant_content(query, k=3) if not context: return "🔍 No relevant information found. Try rephrasing your question." # Use appropriate model for answering if self.qa_pipeline is None: return self.extract_direct_answer(query, context) try: if self.model_type == "distilbert-qa" or self.model_type == "fallback": # Use Q&A pipeline result = self.qa_pipeline(question=query, context=context) answer = result['answer'] confidence = result['score'] if confidence > 0.1: # Reasonable confidence return f"**Answer:** {answer}\n\n**Context:** {context[:200]}..." else: return self.extract_direct_answer(query, context) elif self.model_type == "flan-t5": # Use text generation pipeline prompt = f"Answer the question based on the context.\nContext: {context}\nQuestion: {query}\nAnswer:" result = self.qa_pipeline(prompt, max_length=200, num_return_sequences=1) answer = result[0]['generated_text'].replace(prompt, '').strip() return f"**Answer:** {answer}" else: return self.extract_direct_answer(query, context) except Exception as e: print(f"Model inference error: {e}") return self.extract_direct_answer(query, context) except Exception as e: return f"❌ Error processing question: {str(e)}" def extract_direct_answer(self, query: str, context: str) -> str: """Direct answer extraction as fallback""" query_lower = query.lower() # Name extraction if any(word in query_lower for word in ['name', 'who is', 'who']): names = re.findall(r'\b[A-Z][a-z]+ [A-Z][a-z]+\b', context) if names: return f"**Name:** {names[0]}" # Experience extraction if any(word in query_lower for word in ['experience', 'years']): exp_matches = re.findall(r'(\d+)[\+\-\s]*(?:years?|yrs?)', context.lower()) if exp_matches: return f"**Experience:** {exp_matches[0]} years" # Skills extraction if any(word in query_lower for word in ['skill', 'technology', 'tech']): # Common tech skills tech_patterns = [ r'\b(?:Python|Java|JavaScript|React|Node|SQL|AWS|Docker|Kubernetes|Git)\b', r'\b(?:HTML|CSS|Angular|Vue|Spring|Django|Flask|MongoDB|PostgreSQL)\b' ] skills = [] for pattern in tech_patterns: skills.extend(re.findall(pattern, context, re.IGNORECASE)) if skills: return f"**Skills mentioned:** {', '.join(set(skills))}" # Education extraction if any(word in query_lower for word in ['education', 'degree', 'university']): edu_matches = re.findall(r'(?:Bachelor|Master|PhD|B\.?S\.?|M\.?S\.?|B\.?A\.?|M\.?A\.?).*?(?:in|of)\s+([^.]+)', context) if edu_matches: return f"**Education:** {edu_matches[0]}" # Return first relevant sentence sentences = [s.strip() for s in context.split('.') if s.strip()] if sentences: return f"**Answer:** {sentences[0]}" return "I found relevant content but couldn't extract a specific answer." # Initialize the system print("Initializing Enhanced Smart RAG System...") rag_system = SmartDocumentRAG() # Create the interface def create_interface(): with gr.Blocks(title="🧠 Enhanced Document Q&A", theme=gr.themes.Soft()) as demo: gr.Markdown(""" # 🧠 Enhanced Document Q&A System **Optimized with Better Models & Quantization!** **Features:** - 🎯 Flan-T5 or DistilBERT for accurate Q&A - ⚡ 4-bit quantization for GPU efficiency - 📊 Direct answer extraction - 🔍 Enhanced semantic search """) with gr.Tab("📤 Upload & Process"): with gr.Row(): with gr.Column(): file_upload = gr.File( label="📁 Upload Documents", file_count="multiple", file_types=[".pdf", ".docx", ".txt"], height=150 ) process_btn = gr.Button("🔄 Process Documents", variant="primary", size="lg") with gr.Column(): process_status = gr.Textbox( label="📋 Processing Status", lines=10, interactive=False ) process_btn.click( fn=rag_system.process_documents, inputs=[file_upload], outputs=[process_status] ) with gr.Tab("❓ Q&A"): with gr.Row(): with gr.Column(): question_input = gr.Textbox( label="🤔 Ask Your Question", placeholder="What is the person's name? / How many years of experience? / What skills do they have?", lines=3 ) with gr.Row(): ask_btn = gr.Button("🧠 Get Answer", variant="primary") summary_btn = gr.Button("📊 Get Summary", variant="secondary") with gr.Column(): answer_output = gr.Textbox( label="💡 Answer", lines=8, interactive=False ) ask_btn.click( fn=rag_system.answer_question, inputs=[question_input], outputs=[answer_output] ) summary_btn.click( fn=lambda: rag_system.answer_question("summary"), inputs=[], outputs=[answer_output] ) return demo # Launch the app if __name__ == "__main__": demo = create_interface() demo.launch( server_name="0.0.0.0", server_port=7860, share=True )