hogragger

hogragger.py

@@ -0,0 +1,251 @@
+import json
+import re
+import nltk
+from nltk.tokenize import sent_tokenize
+import torch
+from sentence_transformers import SentenceTransformer, util
+import faiss
+import numpy as np
+from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
+from rank_bm25 import BM25Okapi  # BM25 for hybrid search
+import logging
+
+
+nltk.download('punkt', quiet=True)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+
+class Hogragger:
+    def __init__(self, corpus_path, model_name='sentence-transformers/all-MiniLM-L12-v2', qa_model='deepset/roberta-large-squad2', classifier_model='deepset/roberta-large-squad2'):
+        self.corpus = self.load_corpus(corpus_path)
+        self.cleaned_passages = self.preprocess_corpus()
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        logging.info(f"Using device: {self.device}")
+
+        # Initialize embedding model and build FAISS index
+        self.model = SentenceTransformer(model_name).to(self.device)
+        self.index = self.build_faiss_index()
+
+        # Initialize BM25 for lexical matching
+        self.bm25 = self.build_bm25_index()
+
+        # Initialize classifier for question type prediction
+        self.tokenizer = AutoTokenizer.from_pretrained(classifier_model)
+        self.classifier = AutoModelForSequenceClassification.from_pretrained(classifier_model).to(self.device)
+
+        # QA Model
+        self.qa_model = pipeline('question-answering', model=qa_model, device=0 if self.device == 'cuda' else -1)
+
+    def load_corpus(self, path):
+        logging.info(f"Loading corpus from {path}")
+        with open(path, "r") as f:
+            corpus = json.load(f)
+        logging.info(f"Loaded {len(corpus)} documents")
+        return corpus
+
+    # def preprocess_corpus(self):
+    #     cleaned_passages = []
+    #     for article in self.corpus:
+    #         body = article.get('body', '')
+    #         clean_body = re.sub(r'<.*?>', '', body)  # Clean HTML tags
+    #         clean_body = re.sub(r'\s+', ' ', clean_body).strip()  # Clean extra spaces
+    #         sentences = sent_tokenize(clean_body)
+
+    #         chunk = ""
+    #         for sentence in sentences:
+    #             if len(chunk.split()) + len(sentence.split()) <= 300:
+    #                 chunk += " " + sentence
+    #             else:
+    #                 cleaned_passages.append(self.create_passage(article, chunk))
+    #                 chunk = sentence
+
+    #         if chunk:
+    #             cleaned_passages.append(self.create_passage(article, chunk))
+    #     logging.info(f"Created {len(cleaned_passages)} passages")
+    #     return cleaned_passages
+    def preprocess_corpus(self):
+        cleaned_passages = []
+        for article in self.corpus:
+            body = article.get('body', '')
+            clean_body = re.sub(r'<.*?>', '', body)  # Clean HTML tags
+            clean_body = re.sub(r'\s+', ' ', clean_body).strip()  # Clean extra spaces
+
+            # Simply take the full cleaned text as a passage without chunking or sentence splitting
+            cleaned_passages.append(self.create_passage(article, clean_body))
+
+        logging.info(f"Created {len(cleaned_passages)} passages")
+        return cleaned_passages
+
+    def create_passage(self, article, chunk):
+        """Creates a passage dictionary from an article and chunk of text."""
+        return {
+            "title": article['title'],
+            "author": article.get('author', 'Unknown'),
+            "published_at": article['published_at'],
+            "category": article['category'],
+            "url": article['url'],
+            "source": article['source'],
+            "passage": chunk.strip()
+        }
+
+    def build_faiss_index(self):
+        logging.info("Building FAISS index...")
+        embeddings = self.model.encode([p['passage'] for p in self.cleaned_passages], convert_to_tensor=True, device=self.device)
+        embeddings = np.array(embeddings.cpu()).astype('float32')
+        logging.info(f"Shape of embeddings: {embeddings.shape}")
+
+        index = faiss.IndexFlatL2(embeddings.shape[1])  # Initialize FAISS index
+
+        if self.device == 'cuda':
+            try:
+                res = faiss.StandardGpuResources()
+                gpu_index = faiss.index_cpu_to_gpu(res, 0, index)
+                gpu_index.add(embeddings)
+                logging.info("Successfully created GPU index")
+                return gpu_index
+            except RuntimeError as e:
+                logging.error(f"GPU index creation failed: {e}")
+                logging.info("Falling back to CPU index")
+
+        index.add(embeddings)  # Add embeddings to CPU index
+        logging.info("Successfully created CPU index")
+        return index
+
+    def build_bm25_index(self):
+        logging.info("Building BM25 index...")
+        tokenized_corpus = [p['passage'].split() for p in self.cleaned_passages]
+        bm25 = BM25Okapi(tokenized_corpus)
+        logging.info("Successfully built BM25 index")
+        return bm25
+
+    def predict_question_type(self, query):
+        inputs = self.tokenizer(query, return_tensors='pt').to(self.device)
+        outputs = self.classifier(**inputs)
+        prediction = torch.argmax(outputs.logits, dim=1).item()
+
+        labels = {0: 'inference_query', 1: 'comparison_query', 2: 'null_query', 3: 'temporal_query', 4: 'fact_query'}
+        return labels.get(prediction, 'unknown_query')
+
+    def retrieve_passages(self, query, k=100, threshold=0.7):
+        try:
+            # FAISS retrieval
+            query_embedding = self.model.encode([query], convert_to_tensor=True, device=self.device)
+            D, I = self.index.search(np.array(query_embedding.cpu()), k)
+
+            # BM25 retrieval
+            tokenized_query = query.split()
+            bm25_scores = self.bm25.get_scores(tokenized_query)
+
+            # Combine FAISS and BM25 results
+            hybrid_scores = self.combine_faiss_bm25_scores(D[0], bm25_scores, I)
+
+            # Filter passages based on hybrid score
+            passages = [self.cleaned_passages[i] for i, score in zip(I[0], hybrid_scores) if score > threshold]
+
+            logging.info(f"Retrieved {len(passages)} passages using hybrid search for query.")
+            return passages
+        except Exception as e:
+            logging.error(f"Error in retrieving passages: {e}")
+            return []
+
+    def combine_faiss_bm25_scores(self, faiss_scores, bm25_scores, passage_indices):
+        # Normalize and combine FAISS and BM25 scores
+        bm25_scores = np.array(bm25_scores)[passage_indices]
+        faiss_scores = np.array(faiss_scores)
+
+        # Convert FAISS distances into similarities by inverting the scale
+        faiss_similarities = 1 / (faiss_scores + 1e-6)  # Avoid division by zero
+
+        # Normalize scores (scale between 0 and 1)
+        bm25_scores = (bm25_scores - np.min(bm25_scores)) / (np.max(bm25_scores) - np.min(bm25_scores) + 1e-6)
+        faiss_similarities = (faiss_similarities - np.min(faiss_similarities)) / (np.max(faiss_similarities) - np.min(faiss_similarities) + 1e-6)
+
+        # Weighted combination (you can adjust weights)
+        combined_scores = 0.7 * faiss_similarities + 0.3 * bm25_scores
+        combined_scores = np.squeeze(combined_scores)  # Ensure it's a single-dimensional array
+
+        return combined_scores
+
+    def filter_passages(self, query, passages):
+        try:
+            query_embedding = self.model.encode(query, convert_to_tensor=True)
+            passage_embeddings = self.model.encode([p['passage'] for p in passages], convert_to_tensor=True)
+
+            similarities = util.pytorch_cos_sim(query_embedding, passage_embeddings)
+            top_k = min(10, len(passages))
+            top_indices = similarities.topk(k=top_k)[1].tolist()[0]
+
+            selected_passages = []
+            used_titles = set()
+            for i in top_indices:
+                if passages[i]['title'] not in used_titles:
+                    selected_passages.append(passages[i])
+                    used_titles.add(passages[i]['title'])
+
+            return selected_passages
+        except Exception as e:
+            logging.error(f"Error in filtering passages: {e}")
+            return []
+
+    def generate_answer(self, query, passages):
+        try:
+            context = " ".join([p['passage'] for p in passages[:5]])
+            answer = self.qa_model(question=query, context=context)
+            logging.info(f"Generated answer: {answer['answer']}")
+            return answer['answer']
+        except Exception as e:
+            logging.error(f"Error in generating answer: {e}")
+            return "Insufficient information."
+
+    def post_process_answer(self, answer, confidence=0.2):
+        answer = re.sub(r'^.*\?', '', answer).strip()
+        answer = answer.capitalize()
+
+        if len(answer) > 100:
+            truncated = re.match(r'^(.*?[.!?])\s', answer)
+            if truncated:
+                answer = truncated.group(1)
+
+        if confidence < 0.2:
+            logging.warning(f"Answer confidence too low: {confidence}")
+            return "I'm unsure about this answer."
+
+        return answer
+
+    def process_query(self, query):
+        question_type = self.predict_question_type(query)
+        retrieved_passages = self.retrieve_passages(query, k=100, threshold=0.7)
+        if not retrieved_passages:
+            return {"query": query, "answer": "No relevant information found", "question_type": question_type, "evidence_list": []}
+
+        filtered_passages = self.filter_passages(query, retrieved_passages)
+        raw_answer = self.generate_answer(query, filtered_passages)
+
+        evidence_count = min(len(filtered_passages), 4)
+        evidence_list = [
+            {
+                "title": p['title'],
+                "author": p['author'],
+                "url": p['url'],
+                "source": p['source'],
+                "category": p['category'],
+                "published_at": p['published_at'],
+                "fact": self.extract_fact(p['passage'], query)
+            } for p in filtered_passages[:evidence_count]
+        ]
+        final_answer = self.post_process_answer(raw_answer)
+
+        return {
+            "query": query,
+            "answer": final_answer,
+            "question_type": question_type,
+            "evidence_list": evidence_list
+        }
+
+    def extract_fact(self, passage, query):
+        # Extracting most relevant sentence from passage
+        sentences = sent_tokenize(passage)
+        query_keywords = set(query.lower().split())
+
+        best_sentence = max(sentences, key=lambda s: len(set(s.lower().split()) & query_keywords), default="")
+        return best_sentence if best_sentence else (sentences[0] if sentences else "")