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