app.py

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, max_chunk_size: int = 300, overlap: int = 50) -> list[str]:
        """
        Splits text into smaller overlapping chunks for better semantic search.
        
        Args:
            text (str): The full text to chunk.
            max_chunk_size (int): Maximum tokens/words per chunk.
            overlap (int): Number of words overlapping between consecutive chunks.
    
        Returns:
            list[str]: List of text chunks.
        """
        import re
    
        # Clean and normalize whitespace
        text = re.sub(r'\s+', ' ', text).strip()
        
        words = text.split()
        chunks = []
        start = 0
        text_len = len(words)
    
        while start < text_len:
            end = min(start + max_chunk_size, text_len)
            chunk_words = words[start:end]
            chunk = ' '.join(chunk_words)
            chunks.append(chunk)
            # Move start forward by chunk size minus overlap to create overlap
            start += max_chunk_size - overlap
    
        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 with stricter relevance filter"""
        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]):
                score = scores[0][i]
                if idx < len(self.documents) and score > 0.4:  # ✅ stricter similarity filter
                    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 and hallucination reduction."""
        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 explicitly
            if any(word in query_lower for word in ['summary', 'summarize', 'about', 'overview']):
                return f"📄 **Document Summary:**\n\n{self.document_summary}"
            
            # Retrieve relevant content chunks via semantic search
            context = self.find_relevant_content(query, k=3)
            
            if not context:
                return "🔍 No relevant information found. Try rephrasing your question."
            
            # If no QA pipeline, fall back to direct extraction
            if self.qa_pipeline is None:
                return self.extract_direct_answer(query, context)
            
            try:
                if self.model_type in ["distilbert-qa", "fallback"]:
                    # Use extractive Q&A pipeline
                    result = self.qa_pipeline(question=query, context=context)
                    answer = result.get('answer', '').strip()
                    confidence = result.get('score', 0)
                    
                    if confidence > 0.1 and answer:
                        return f"**Answer:** {answer}\n\n**Context:** {context[:200]}..."
                    else:
                        return self.extract_direct_answer(query, context)
                
                elif self.model_type == "flan-t5":
                    # Use generative model with improved prompt to reduce hallucination
                    prompt = (
                        f"Answer concisely and strictly based on the following context.\n\n"
                        f"Context:\n{context}\n\n"
                        f"Question:\n{query}\n\n"
                        f"If the answer is not contained in the context, reply with 'Not found in document.'\n"
                        f"Answer:"
                    )
                    result = self.qa_pipeline(prompt, max_length=256, num_return_sequences=1)
                    generated_text = result[0].get('generated_text', '')
                    answer = generated_text.replace(prompt, '').strip()
                    
                    if answer.lower() in ["not found in document.", "no answer", "unknown", ""]:
                        return "🔍 Sorry, the answer was not found in the documents."
                    else:
                        return f"**Answer:** {answer}"
                
                else:
                    # Default fallback extraction
                    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."

    def clean_text(self, text: str) -> str:
        """
        Clean and normalize raw text by:
        - Removing excessive whitespace
        - Fixing merged words (camel case separation)
        - Removing unwanted characters (optional)
        - Lowercasing or preserving case (optional)
        """
        import re
    
        # Replace multiple whitespace/newlines/tabs with single space
        text = re.sub(r'\s+', ' ', text).strip()
        
        # Fix merged words like 'wordAnotherWord' -> 'word Another Word'
        text = re.sub(r'([a-z])([A-Z])', r'\1 \2', text)
        
        # Optional: remove special characters except basic punctuation
        # text = re.sub(r'[^a-zA-Z0-9,.!?;:\'\"()\-\s]', '', text)
        
        return text





    

# 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
    )