app.py

'''

import os
import re
import json
import torch
import numpy as np
import logging
from typing import Dict, List, Tuple, Optional
from tqdm import tqdm
from pydantic import BaseModel
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    AutoModelForQuestionAnswering,
    pipeline,
    LogitsProcessor,
    LogitsProcessorList,
    PreTrainedModel,
    PreTrainedTokenizer
)
from sentence_transformers import SentenceTransformer, CrossEncoder
from sklearn.feature_extraction.text import TfidfVectorizer
from rank_bm25 import BM25Okapi
import PyPDF2
from sklearn.cluster import KMeans
import spacy
import subprocess
import gradio as gr

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s"
)

class ConfidenceCalibrator(LogitsProcessor):
    def __init__(self, calibration_factor: float = 0.9):
        self.calibration_factor = calibration_factor

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        return scores / self.calibration_factor

class DocumentResult(BaseModel):
    content: str
    confidence: float
    source_page: int
    supporting_evidence: List[str]

class OptimalModelSelector:
    def __init__(self):
        self.qa_models = {
            "deberta-v3": ("deepset/deberta-v3-large-squad2", 0.87)
        }
        self.summarization_models = {
            "bart": ("facebook/bart-large-cnn", 0.85)
        }
        self.current_models = {}

    def get_best_model(self, task_type: str) -> Tuple[PreTrainedModel, PreTrainedTokenizer, float]:
        model_map = self.qa_models if "qa" in task_type else self.summarization_models
        best_model_name, best_score = max(model_map.items(), key=lambda x: x[1][1])
        if best_model_name not in self.current_models:
            tokenizer = AutoTokenizer.from_pretrained(model_map[best_model_name][0])
            model = (AutoModelForQuestionAnswering if "qa" in task_type
                     else AutoModelForSeq2SeqLM).from_pretrained(model_map[best_model_name][0])
            model = model.eval().half().to('cuda' if torch.cuda.is_available() else 'cpu')
            self.current_models[best_model_name] = (model, tokenizer)
        return *self.current_models[best_model_name], best_score

class PDFAugmentedRetriever:
    def __init__(self, document_texts: List[str]):
        self.documents = [(i, text) for i, text in enumerate(document_texts)]
        self.bm25 = BM25Okapi([text.split() for _, text in self.documents])
        self.encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
        self.tfidf = TfidfVectorizer(stop_words='english').fit([text for _, text in self.documents])

    def retrieve(self, query: str, top_k: int = 5) -> List[Tuple[int, str, float]]:
        bm25_scores = self.bm25.get_scores(query.split())
        semantic_scores = self.encoder.predict([(query, doc) for _, doc in self.documents])
        combined_scores = 0.4 * bm25_scores + 0.6 * np.array(semantic_scores)
        top_indices = np.argsort(combined_scores)[-top_k:][::-1]
        return [(self.documents[i][0], self.documents[i][1], float(combined_scores[i]))
                for i in top_indices]

class DetailedExplainer:
    def __init__(self,
                 explanation_model: str = "google/flan-t5-large",
                 device: int = 0):
        try:
            self.nlp = spacy.load("en_core_web_sm")
        except OSError:
            subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"], check=True)
            self.nlp = spacy.load("en_core_web_sm")
        self.explainer = pipeline(
            "text2text-generation",
            model=explanation_model,
            tokenizer=explanation_model,
            device=device,
            max_length=500,
            max_new_tokens=800
        )

    def extract_concepts(self, text: str) -> list:
        doc = self.nlp(text)
        concepts = set()
        for chunk in doc.noun_chunks:
            if len(chunk) > 1 and not chunk.root.is_stop:
                concepts.add(chunk.text.strip())
        for ent in doc.ents:
            if ent.label_ in ["PERSON", "ORG", "GPE", "NORP", "EVENT", "WORK_OF_ART"]:
                concepts.add(ent.text.strip())
        return list(concepts)

    def explain_concept(self, concept: str, context: str, min_accuracy: float = 0.50) -> str:
        prompt = (
            f"The following sentence from a PDF is given \n{context}\n\n\nNow explain the concept '{concept}' mentioned above with at least {int(min_accuracy * 100)}% accuracy."
        )
        result = self.explainer(
            prompt,
            do_sample=False
        )
        return result[0]["generated_text"].strip()

    def explain_text(self, text: str, context: str) -> dict:
        concepts = self.extract_concepts(text)
        explanations = {}
        for concept in concepts:
            explanations[concept] = self.explain_concept(concept, context)
        return {"concepts": concepts, "explanations": explanations}

class AdvancedPDFAnalyzer:
    def __init__(self):
        self.logger = logging.getLogger("PDFAnalyzer")
        self.model_selector = OptimalModelSelector()
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.qa_model, self.qa_tokenizer, _ = self.model_selector.get_best_model("qa")
        self.qa_model = self.qa_model.to(self.device)
        self.summarizer = pipeline(
            "summarization",
            model="facebook/bart-large-cnn",
            device=0 if torch.cuda.is_available() else -1,
            framework="pt"
        )
        self.logits_processor = LogitsProcessorList([
            ConfidenceCalibrator(calibration_factor=0.85)
        ])
        self.detailed_explainer = DetailedExplainer(device=0 if torch.cuda.is_available() else -1)

    def extract_text_with_metadata(self, file_path: str) -> List[Dict]:
        documents = []
        with open(file_path, 'rb') as f:
            reader = PyPDF2.PdfReader(f)
            for i, page in enumerate(reader.pages):
                text = page.extract_text()
                if not text or not text.strip():
                    continue
                page_number = i + 1
                metadata = {
                    'source': os.path.basename(file_path),
                    'page': page_number,
                    'char_count': len(text),
                    'word_count': len(text.split()),
                }
                documents.append({
                    'content': self._clean_text(text),
                    'metadata': metadata
                })
        if not documents:
            raise ValueError("No extractable content found in PDF")
        return documents

    def _clean_text(self, text: str) -> str:
        text = re.sub(r'[\x00-\x1F\x7F-\x9F]', ' ', text)
        text = re.sub(r'\s+', ' ', text)
        text = re.sub(r'(\w)-\s+(\w)', r'\1\2', text)
        return text.strip()

    def answer_question(self, question: str, documents: List[Dict]) -> Dict:
        retriever = PDFAugmentedRetriever([doc['content'] for doc in documents])
        relevant_contexts = retriever.retrieve(question, top_k=3)
        answers = []
        for page_idx, context, similarity_score in relevant_contexts:
            inputs = self.qa_tokenizer(
                question,
                context,
                add_special_tokens=True,
                return_tensors="pt",
                max_length=512,
                truncation=True
            )
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            with torch.no_grad():
                outputs = self.qa_model(**inputs)
                start_logits = outputs.start_logits
                end_logits = outputs.end_logits
                logits_processor = LogitsProcessorList([ConfidenceCalibrator()])
                start_logits = logits_processor(inputs['input_ids'], start_logits)
                end_logits = logits_processor(inputs['input_ids'], end_logits)
                start_prob = torch.nn.functional.softmax(start_logits, dim=-1)
                end_prob = torch.nn.functional.softmax(end_logits, dim=-1)
                max_start_score, max_start_idx = torch.max(start_prob, dim=-1)
                max_start_idx_int = max_start_idx.item()
                max_end_score, max_end_idx = torch.max(end_prob[0, max_start_idx_int:], dim=-1)
                max_end_idx_int = max_end_idx.item() + max_start_idx_int
                confidence = float((max_start_score * max_end_score) * 0.9 * similarity_score)
                answer_tokens = inputs["input_ids"][0][max_start_idx_int:max_end_idx_int + 1]
                answer = self.qa_tokenizer.decode(answer_tokens, skip_special_tokens=True)
                explanations_result = self.detailed_explainer.explain_text(answer, context)
                answers.append({
                    "answer": answer,
                    "confidence": confidence,
                    "context": context,
                    "page_number": documents[page_idx]['metadata']['page'],
                    "explanations": explanations_result
                })
        if not answers:
            return {"answer": "No confident answer found", "confidence": 0.0, "explanations": {}}
        best_answer = max(answers, key=lambda x: x['confidence'])
        if best_answer['confidence'] < 0.85:
            best_answer['answer'] = f"[Low Confidence] {best_answer['answer']}"
            return answers #MODUSTIFIED HERE YOU REMOVE THIS THIS LINE OF CODE IF IT CRASHES, DAT 10TH AUG, 11:49AM
        return best_answer

analyzer = AdvancedPDFAnalyzer()
documents = analyzer.extract_text_with_metadata("example.pdf")

def ask_question_gradio(question: str):
    if not question.strip():
        return "Please enter a valid question."
    try:
        result = analyzer.answer_question(question, documents)
        answer = result['answer']
        confidence = result['confidence']
        explanation = "\n\n".join(
            f"🔹 {concept}: {desc}"
            for concept, desc in result.get("explanations", {}).get("explanations", {}).items()
        )
        return f"📌 **Answer**: {answer}\n\n🔒 **Confidence**: {confidence:.2f}\n\n📘 **Explanations**:\n{explanation}"
    except Exception as e:
        return f"❌ Error: {str(e)}"

demo = gr.Interface(
    fn=ask_question_gradio,
    inputs=gr.Textbox(label="Ask a question about the PDF"),
    outputs=gr.Markdown(label="Answer"),
    title="Quandans AI - Ask Questions",
    description="Ask a question based on the document loaded in this system."
)

demo.launch()
'''


import os
import re
import json
import torch
import numpy as np
import logging
from typing import Dict, List, Tuple, Optional
from tqdm import tqdm
from pydantic import BaseModel
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    AutoModelForQuestionAnswering,
    pipeline,
    LogitsProcessor,
    LogitsProcessorList,
    PreTrainedModel,
    PreTrainedTokenizer
)
from sentence_transformers import SentenceTransformer, CrossEncoder
from sklearn.feature_extraction.text import TfidfVectorizer
from rank_bm25 import BM25Okapi
import PyPDF2
from sklearn.cluster import KMeans
import spacy
import subprocess
import gradio as gr

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s"
)

class ConfidenceCalibrator(LogitsProcessor):
    def __init__(self, calibration_factor: float = 0.9):
        self.calibration_factor = calibration_factor

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        return scores / self.calibration_factor

class DocumentResult(BaseModel):
    content: str
    confidence: float
    source_page: int
    supporting_evidence: List[str]

class OptimalModelSelector:
    def __init__(self):
        self.qa_models = {
            "deberta-v3": ("deepset/deberta-v3-large-squad2", 0.87)
        }
        self.summarization_models = {
            "bart": ("facebook/bart-large-cnn", 0.85)
        }
        self.current_models = {}

    def get_best_model(self, task_type: str) -> Tuple[PreTrainedModel, PreTrainedTokenizer, float]:
        model_map = self.qa_models if "qa" in task_type else self.summarization_models
        best_model_name, best_score = max(model_map.items(), key=lambda x: x[1][1])
        if best_model_name not in self.current_models:
            tokenizer = AutoTokenizer.from_pretrained(model_map[best_model_name][0])
            model = (AutoModelForQuestionAnswering if "qa" in task_type
                     else AutoModelForSeq2SeqLM).from_pretrained(model_map[best_model_name][0])
            model = model.eval().half().to('cuda' if torch.cuda.is_available() else 'cpu')
            self.current_models[best_model_name] = (model, tokenizer)
        return *self.current_models[best_model_name], best_score

class PDFAugmentedRetriever:
    def __init__(self, document_texts: List[str]):
        self.documents = [(i, text) for i, text in enumerate(document_texts)]
        self.bm25 = BM25Okapi([text.split() for _, text in self.documents])
        self.encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
        self.tfidf = TfidfVectorizer(stop_words='english').fit([text for _, text in self.documents])

    def retrieve(self, query: str, top_k: int = 5) -> List[Tuple[int, str, float]]:
        bm25_scores = self.bm25.get_scores(query.split())
        semantic_scores = self.encoder.predict([(query, doc) for _, doc in self.documents])
        combined_scores = 0.4 * bm25_scores + 0.6 * np.array(semantic_scores)
        top_indices = np.argsort(combined_scores)[-top_k:][::-1]
        return [(self.documents[i][0], self.documents[i][1], float(combined_scores[i]))
                for i in top_indices]

class DetailedExplainer:
    def __init__(self,
                 explanation_model: str = "google/flan-t5-large",
                 device: int = 0):
        try:
            self.nlp = spacy.load("en_core_web_sm")
        except OSError:
            subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"], check=True)
            self.nlp = spacy.load("en_core_web_sm")
        self.explainer = pipeline(
            "text2text-generation",
            model=explanation_model,
            tokenizer=explanation_model,
            device=device,
            max_length=500,
            max_new_tokens=800
        )

    def extract_concepts(self, text: str) -> list:
        doc = self.nlp(text)
        concepts = set()
        for chunk in doc.noun_chunks:
            if len(chunk) > 1 and not chunk.root.is_stop:
                concepts.add(chunk.text.strip())
        for ent in doc.ents:
            if ent.label_ in ["PERSON", "ORG", "GPE", "NORP", "EVENT", "WORK_OF_ART"]:
                concepts.add(ent.text.strip())
        return list(concepts)

    def explain_concept(self, concept: str, context: str, min_accuracy: float = 0.50) -> str:
        prompt = (
            f"The following sentence from a PDF is given \n{context}\n\n\nNow explain the concept '{concept}' mentioned above with at least {int(min_accuracy * 100)}% accuracy."
        )
        result = self.explainer(
            prompt,
            do_sample=False
        )
        return result[0]["generated_text"].strip()

    def explain_text(self, text: str, context: str) -> dict:
        concepts = self.extract_concepts(text)
        explanations = {}
        for concept in concepts:
            explanations[concept] = self.explain_concept(concept, context)
        return {"concepts": concepts, "explanations": explanations}

class AdvancedPDFAnalyzer:
    def __init__(self):
        self.logger = logging.getLogger("PDFAnalyzer")
        self.model_selector = OptimalModelSelector()
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.qa_model, self.qa_tokenizer, _ = self.model_selector.get_best_model("qa")
        self.qa_model = self.qa_model.to(self.device)
        self.summarizer = pipeline(
            "summarization",
            model="facebook/bart-large-cnn",
            device=0 if torch.cuda.is_available() else -1,
            framework="pt"
        )
        self.logits_processor = LogitsProcessorList([
            ConfidenceCalibrator(calibration_factor=0.85)
        ])
        self.detailed_explainer = DetailedExplainer(device=0 if torch.cuda.is_available() else -1)

    def extract_text_with_metadata(self, file_path: str) -> List[Dict]:
        documents = []
        with open(file_path, 'rb') as f:
            reader = PyPDF2.PdfReader(f)
            for i, page in enumerate(reader.pages):
                text = page.extract_text()
                if not text or not text.strip():
                    continue
                page_number = i + 1
                metadata = {
                    'source': os.path.basename(file_path),
                    'page': page_number,
                    'char_count': len(text),
                    'word_count': len(text.split()),
                }
                documents.append({
                    'content': self._clean_text(text),
                    'metadata': metadata
                })
        if not documents:
            raise ValueError("No extractable content found in PDF")
        return documents

    def _clean_text(self, text: str) -> str:
        text = re.sub(r'[\x00-\x1F\x7F-\x9F]', ' ', text)
        text = re.sub(r'\s+', ' ', text)
        text = re.sub(r'(\w)-\s+(\w)', r'\1\2', text)
        return text.strip()

    def answer_question(self, question: str, documents: List[Dict]) -> Dict:
        retriever = PDFAugmentedRetriever([doc['content'] for doc in documents])
        relevant_contexts = retriever.retrieve(question, top_k=3)
        answers = []
        
        for page_idx, context, similarity_score in relevant_contexts:
            inputs = self.qa_tokenizer(
                question,
                context,
                add_special_tokens=True,
                return_tensors="pt",
                max_length=512,
                truncation=True
            )
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            
            with torch.no_grad():
                outputs = self.qa_model(**inputs)
                start_logits = outputs.start_logits
                end_logits = outputs.end_logits
                
                logits_processor = LogitsProcessorList([ConfidenceCalibrator()])
                start_logits = logits_processor(inputs['input_ids'], start_logits)
                end_logits = logits_processor(inputs['input_ids'], end_logits)
                
                start_prob = torch.nn.functional.softmax(start_logits, dim=-1)
                end_prob = torch.nn.functional.softmax(end_logits, dim=-1)
                
                max_start_score, max_start_idx = torch.max(start_prob, dim=-1)
                max_start_idx_int = max_start_idx.item()
                max_end_score, max_end_idx = torch.max(end_prob[0, max_start_idx_int:], dim=-1)
                max_end_idx_int = max_end_idx.item() + max_start_idx_int
                
                confidence = float((max_start_score * max_end_score) * 0.9 * similarity_score)
                answer_tokens = inputs["input_ids"][0][max_start_idx_int:max_end_idx_int + 1]
                answer = self.qa_tokenizer.decode(answer_tokens, skip_special_tokens=True)
                
                # Only generate explanations if we have a valid answer
                explanations_result = {"concepts": [], "explanations": {}}
                if answer and answer.strip():
                    try:
                        explanations_result = self.detailed_explainer.explain_text(answer, context)
                    except Exception as e:
                        self.logger.warning(f"Failed to generate explanations: {e}")
                
                answers.append({
                    "answer": answer,
                    "confidence": confidence,
                    "context": context,
                    "page_number": documents[page_idx]['metadata']['page'],
                    "explanations": explanations_result
                })
        
        if not answers:
            return {
                "answer": "No confident answer found", 
                "confidence": 0.0, 
                "explanations": {"concepts": [], "explanations": {}},
                "page_number": 0,
                "context": ""
            }
        
        # Get the best answer based on confidence
        best_answer = max(answers, key=lambda x: x['confidence'])
        
        # FIXED: Always return the best answer dictionary, just modify the answer text if confidence is low
        if best_answer['confidence'] < 0.3:  # Lowered threshold to be more permissive
            best_answer['answer'] = f"[Low Confidence] {best_answer['answer']}"
        
        return best_answer

# Initialize analyzer (make sure to update the PDF path)
analyzer = AdvancedPDFAnalyzer()

# Global variable to store documents
documents = []

def load_pdf(file_path: str):
    """Load PDF and extract documents"""
    global documents
    try:
        documents = analyzer.extract_text_with_metadata(file_path)
        return f"Successfully loaded PDF with {len(documents)} pages."
    except Exception as e:
        return f"Error loading PDF: {str(e)}"

def ask_question_gradio(question: str):
    if not question.strip():
        return "Please enter a valid question."
    
    if not documents:
        return "❌ No PDF loaded. Please load a PDF first."
    
    try:
        result = analyzer.answer_question(question, documents)
        
        # Ensure we have the expected structure
        answer = result.get('answer', 'No answer found')
        confidence = result.get('confidence', 0.0)
        page_number = result.get('page_number', 0)
        explanations = result.get("explanations", {}).get("explanations", {})
        
        # Format explanations
        explanation_text = ""
        if explanations:
            explanation_text = "\n\n".join(
                f"🔹 **{concept}**: {desc}"
                for concept, desc in explanations.items()
                if desc and desc.strip()
            )
        
        # Build response
        response_parts = [
            f"📌 **Answer**: {answer}",
            f"🔒 **Confidence**: {confidence:.2f}",
            f"📄 **Page**: {page_number}"
        ]
        
        if explanation_text:
            response_parts.append(f"📘 **Explanations**:\n{explanation_text}")
        
        return "\n\n".join(response_parts)
        
    except Exception as e:
        return f"❌ Error: {str(e)}"

# Load your PDF here - update the path to your actual PDF file
pdf_path = "example.pdf"
if os.path.exists(pdf_path):
    load_result = load_pdf(pdf_path)
    print(load_result)
else:
    print(f"PDF file '{pdf_path}' not found. Please update the path.")

demo = gr.Interface(
    fn=ask_question_gradio,
    inputs=gr.Textbox(label="Ask a question about the PDF", placeholder="Type your question here..."),
    outputs=gr.Markdown(label="Answer"),
    title="Quandans AI - Ask Questions",
    description="Ask a question based on the document loaded in this system.",
    examples=[
        "What is the main topic of this document?",
        "Summarize the key points from page 1",
        "What are the conclusions mentioned?"
    ]
)

demo.launch()