Create document_chunker.py

document_chunker.py

@@ -0,0 +1,175 @@
+import re
+from typing import List, Dict, Optional
+from pathlib import Path
+from collections import defaultdict
+from dataclasses import dataclass
+
+from docx import Document
+from sentence_transformers import SentenceTransformer
+from sklearn.feature_extraction.text import TfidfVectorizer
+
+
+@dataclass
+class DocumentChunk:
+    chunk_id: int
+    text: str
+    embedding: List[float]
+    metadata: Dict
+
+
+class DocumentChunker:
+    def __init__(self):
+        self.embed_model = SentenceTransformer("all-MiniLM-L6-v2")
+
+        self.category_patterns = {
+            "Project Summary": [r"\bsummary\b", r"\bproject overview\b"],
+            "Contact Information": [r"\bcontact\b", r"\bemail\b", r"\bphone\b", r"\baddress\b"],
+            "Problem/ Need": [r"\bproblem\b", r"\bneed\b", r"\bchallenge\b"],
+            "Mission Statement": [r"\bmission\b", r"\bvision\b"],
+            "Fit or Alignment to Grant": [r"\balignment\b", r"\bfit\b", r"\bgrant (focus|priority)\b"],
+            "Goals/ Vision / Objectives": [r"\bgoals?\b", r"\bobjectives?\b", r"\bvision\b"],
+            "Our Solution *PROGRAMS* and Approach": [r"\bsolution\b", r"\bprogram\b", r"\bapproach\b"],
+            "Impact, Results, or Outcomes": [r"\bimpact\b", r"\bresults?\b", r"\boutcomes?\b"],
+            "Beneficiaries": [r"\bbeneficiaries\b", r"\bwho we serve\b", r"\btarget audience\b"],
+            "Differentiation with Competitors": [r"\bcompetitor\b", r"\bdifferent\b", r"\bvalue proposition\b"],
+            "Plan and Timeline": [r"\btimeline\b", r"\bschedule\b", r"\bmilestone\b"],
+            "Budget and Funding": [r"\bbudget\b", r"\bfunding\b", r"\bcost\b"],
+            "Sustainability and Strategy": [r"\bsustainability\b", r"\bexit strategy\b"],
+            "Organization's History": [r"\bhistory\b", r"\borganization background\b"],
+            "Team Member Descriptions": [r"\bteam\b", r"\bstaff\b", r"\blived experience\b"],
+        }
+
+        self.patterns = {
+            'grant_application': {
+                'header_patterns': [
+                    r'\*\*([^*]+)\*\*',
+                    r'^([A-Z][^a-z]*[A-Z])$',
+                    r'^([A-Z][A-Za-z\s]+)$',
+                ],
+                'question_patterns': [
+                    r'^.+\?$',
+                    r'^\*?Please .+',
+                    r'^How .+',
+                    r'^What .+',
+                    r'^Describe .+',
+                ]
+            }
+        }
+
+    def match_category(self, text: str, return_first: bool = True) -> Optional[str] or List[str]:
+        lower_text = text.lower()
+        match_scores = defaultdict(int)
+        for category, patterns in self.category_patterns.items():
+            for pattern in patterns:
+                matches = re.findall(pattern, lower_text)
+                match_scores[category] += len(matches)
+
+        if not match_scores:
+            return None if return_first else []
+
+        sorted_categories = sorted(match_scores.items(), key=lambda x: -x[1])
+        return sorted_categories[0][0] if return_first else [cat for cat, _ in sorted_categories if match_scores[cat] > 0]
+
+    def extract_text_from_docx(self, file_path: str) -> str:
+        doc = Document(file_path)
+        return '\n'.join([f"**{p.text}**" if any(r.bold for r in p.runs) else p.text for p in doc.paragraphs])
+
+    def detect_document_type(self, text: str) -> str:
+        keywords = ['grant', 'funding', 'mission']
+        return 'grant_application' if sum(k in text.lower() for k in keywords) >= 2 else 'generic'
+
+    def extract_headers(self, text: str, doc_type: str) -> List[Dict]:
+        lines = text.split('\n')
+        headers = []
+        patterns = self.patterns.get(doc_type, self.patterns['grant_application'])
+        for i, line in enumerate(lines):
+            line = line.strip("* ")
+            if any(re.match(p, line, re.IGNORECASE) for p in patterns['question_patterns']):
+                headers.append({'text': line, 'line_number': i, 'pattern_type': 'question'})
+            elif any(re.match(p, line) for p in patterns['header_patterns']):
+                headers.append({'text': line, 'line_number': i, 'pattern_type': 'header'})
+        return headers
+
+    def chunk_by_headers(self, text: str, headers: List[Dict], max_words=150) -> List[Dict]:
+        lines = text.split('\n')
+        chunks = []
+
+        if not headers:
+            # fallback chunking
+            words = text.split()
+            for i in range(0, len(words), max_words):
+                piece = ' '.join(words[i:i + max_words])
+                chunks.append({
+                    'chunk_id': len(chunks) + 1,
+                    'header': '',
+                    'questions': [],
+                    'content': piece,
+                    'pattern_type': 'auto'
+                })
+            return chunks
+
+        for i, header in enumerate(headers):
+            start, end = header['line_number'], headers[i + 1]['line_number'] if i + 1 < len(headers) else len(lines)
+            content_lines = lines[start + 1:end]
+            questions = [l.strip() for l in content_lines if l.strip().endswith('?') and len(l.split()) <= 20]
+            content = ' '.join([l.strip() for l in content_lines if l.strip() and l.strip() not in questions])
+
+            for j in range(0, len(content.split()), max_words):
+                chunk_text = ' '.join(content.split()[j:j + max_words])
+                chunks.append({
+                    'chunk_id': len(chunks) + 1,
+                    'header': header['text'] if header['pattern_type'] == 'header' else '',
+                    'questions': questions if header['pattern_type'] == 'question' else [],
+                    'content': chunk_text,
+                    'pattern_type': header['pattern_type'],
+                    'split_index': j // max_words
+                })
+        return chunks
+
+    def extract_topics_tfidf(self, text: str, max_features: int = 3) -> List[str]:
+        clean = re.sub(r'[^\w\s]', ' ', text.lower())
+        vectorizer = TfidfVectorizer(max_features=max_features * 2, stop_words='english')
+        tfidf = vectorizer.fit_transform([clean])
+        terms = vectorizer.get_feature_names_out()
+        scores = tfidf.toarray()[0]
+        top_terms = [term for term, score in sorted(zip(terms, scores), key=lambda x: -x[1]) if score > 0]
+        return top_terms[:max_features]
+
+    def calculate_confidence_score(self, chunk: Dict) -> float:
+        score = 0.0
+        if chunk.get('header'): score += 0.3
+        if chunk.get('content') and len(chunk['content'].split()) > 20: score += 0.3
+        if chunk.get('questions'): score += 0.2
+        return min(score, 1.0)
+
+    def process_document(self, file_path: str, title: Optional[str] = None) -> List[Dict]:
+        file_path = Path(file_path)
+        text = self.extract_text_from_docx(str(file_path)) if file_path.suffix == ".docx" else file_path.read_text()
+        doc_type = self.detect_document_type(text)
+        headers = self.extract_headers(text, doc_type)
+        raw_chunks = self.chunk_by_headers(text, headers)
+
+        final_chunks = []
+        for chunk in raw_chunks:
+            full_text = f"{chunk['header']} {' '.join(chunk['questions'])} {chunk['content']}".strip()
+            category = self.match_category(full_text, return_first=True)
+            categories = self.match_category(full_text, return_first=False)
+            embedding = self.embed_model.encode(full_text).tolist()
+            topics = self.extract_topics_tfidf(full_text)
+            confidence = self.calculate_confidence_score(chunk)
+
+            final_chunks.append({
+                "chunk_id": chunk['chunk_id'],
+                "text": full_text,
+                "embedding": embedding,
+                "metadata": {
+                    **chunk,
+                    "title": title or file_path.name,
+                    "category": category,
+                    "categories": categories,
+                    "topics": topics,
+                    "confidence_score": confidence
+                }
+            })
+
+        return final_chunks