malepati/custom_template_working

1import asyncio2from typing import List3 4from models.document_chunk import DocumentChunk5 6 7class ScoreBasedChunker:8 9    def extract_headings(self, text: str) -> List[str]:10        lines = text.split("\n")11        headings = []12        13        for line in lines:14            line = line.strip()15            if line.startswith("#"):16                headings.append(line)17        18        return headings19 20    def score_headings(self, headings: List[str]) -> List[float]:21        heading_scores = []22        last_heading_index = -123        first_heading_found = False24 25        for i, heading in enumerate(headings):26            score = 0.027            28            heading_level = len(heading) - len(heading.lstrip("#"))29            30            if heading_level <= 3:31                score += 10.0 - (heading_level - 1) * 2.032            else:33                score += 4.0 - (heading_level - 4) * 0.534 35            if not first_heading_found:36                score += 5.037                first_heading_found = True38 39            if last_heading_index != -1:40                distance = i - last_heading_index41                distance_bonus = min(5.0, distance * 0.5)42                score += distance_bonus43 44            last_heading_index = i45            heading_scores.append(score)46 47        return heading_scores48 49    def get_chunks_from_headings(50        self,51        text: str,52        headings: List[str],53        heading_scores: List[float],54        top_k: int = 10,55    ) -> List[DocumentChunk]:56        if not heading_scores:57            heading_scores = self.score_headings(headings)58 59        chunks = []60        heading_indices = []61 62        for i, score in enumerate(heading_scores):63            if score > 0:64                heading_indices.append((i, score))65 66        if len(heading_indices) == 0:67            return chunks68 69        heading_indices.sort(key=lambda x: (-x[1], x[0]))70 71        if len(heading_indices) <= top_k:72            selected_indices = [idx for idx, _ in heading_indices]73            selected_indices.sort()74        else:75            score_groups = {}76            for idx, score in heading_indices:77                rounded_score = round(score)78                if rounded_score not in score_groups:79                    score_groups[rounded_score] = []80                score_groups[rounded_score].append(idx)81 82            sorted_groups = sorted(83                score_groups.items(), key=lambda x: x[0], reverse=True84            )85 86            selected_indices = []87 88            for score, indices in sorted_groups:89                indices.sort()90                remaining_needed = top_k - len(selected_indices)91 92                if remaining_needed <= 0:93                    break94 95                if len(indices) <= remaining_needed:96                    selected_indices.extend(indices)97                else:98                    if remaining_needed == 1:99                        mid_idx = len(indices) // 2100                        selected_indices.append(indices[mid_idx])101                    elif remaining_needed == 2:102                        selected_indices.append(indices[0])103                        selected_indices.append(indices[-1])104                    else:105                        step = (len(indices) - 1) / (remaining_needed - 1)106 107                        for i in range(remaining_needed):108                            index = int(round(i * step))109                            if index < len(indices):110                                selected_indices.append(indices[index])111 112            selected_indices.sort()113 114        lines = text.split("\n")115        heading_positions = {}116        117        for i, line in enumerate(lines):118            line_stripped = line.strip()119            if line_stripped.startswith("#"):120                for heading_idx, heading in enumerate(headings):121                    if heading == line_stripped and heading_idx not in heading_positions:122                        heading_positions[heading_idx] = i123                        break124        125        for i, heading_idx in enumerate(selected_indices):126            if heading_idx not in heading_positions:127                continue128                129            heading = headings[heading_idx]130            heading_line_idx = heading_positions[heading_idx]131            132            if i + 1 < len(selected_indices):133                next_heading_idx = selected_indices[i + 1]134                if next_heading_idx in heading_positions:135                    next_heading_line_idx = heading_positions[next_heading_idx]136                    content_end = next_heading_line_idx137                else:138                    content_end = len(lines)139            else:140                content_end = len(lines)141 142            content_lines = lines[heading_line_idx + 1 : content_end]143            content = "\n".join(content_lines).strip()144 145            chunk = DocumentChunk(146                heading=heading,147                content=content,148                heading_index=heading_idx,149                score=heading_scores[heading_idx],150            )151            chunks.append(chunk)152            153        return chunks154 155    async def get_n_chunks(self, text: str, n: int) -> List[DocumentChunk]:156        headings = await asyncio.to_thread(self.extract_headings, text)157        heading_scores = await asyncio.to_thread(self.score_headings, headings)158        chunks = await asyncio.to_thread(159            self.get_chunks_from_headings, text, headings, heading_scores, n160        )161        if len(chunks) < n:162            raise ValueError(f"Only {len(chunks)} chunks found, requested {n}")163        return chunks164