hie-rag 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

hie_rag/hie_rag.py

@@ -15,18 +15,23 @@ class HieRag:
     
     def process_and_save_index_stream(self, file_name: str, uploaded_file: bytes, min_chunk_size, max_chunk_size):
         yield {"status": "🔍 Extracting text..."}
+        print(f"Extracting text from {file_name}")
         extracted_text = self.utils.extract_text(uploaded_file)
 
         yield {"status": "✂️ Splitting into chunks..."}
+        print(f"Splitting text into chunks with min size {min_chunk_size} and max size {max_chunk_size}")
         result_split = self.split.split(extracted_text, min_chunk_size=min_chunk_size, max_chunk_size=max_chunk_size)
 
         yield {"status": "🧠 Processing chunks..."}
+        print(f"Processing {len(result_split)} chunks")
         result_process = self.process.process_chunks(result_split)
 
         yield {"status": "🌲 Building tree index..."}
+        print(f"Building tree index with {len(result_process)} chunks")
         tree_index = self.tree_index.tree_index(file_name = file_name, chunk_metadata=result_process)
 
         yield {"status": "💾 Saving to vector DB..."}
+        print(f"Saving tree index with {len(tree_index.get('chunks', []))} chunks to vector DB")
         save_result = self.vector_db.save_index(tree_index)
 
         file_id = save_result.get("file_id", "unknown")

hie_rag/process.py

@@ -9,8 +9,8 @@ from .utils import Utils
 
 
 class Process:
-    def __init__(self, base_url=None):
-        self.client = ChatOllama(model="llama3.2:latest")
+    def __init__(self, base_url=None, model="llama3.2:latest"):
+        self.client = ChatOllama(model=model)
         self.utils = Utils(base_url=base_url)
 
     def _generate_metadata(self, chunk: str) -> Dict:

hie_rag/split.py

@@ -1,4 +1,6 @@
-from typing import List
+import re
+from collections import deque
+from typing import List, Tuple
 
 from .utils import Utils
 
@@ -10,6 +12,20 @@ class Split:
         """
         self.utils = Utils(base_url=base_url)
 
+    def _custom_split(self, text: str):
+        stripped = text.strip()
+        # 以「空白行」作為段落切點
+        raw_paragraphs = re.split(r'\n\s*\n+', stripped)
+
+        result = []
+        for para in raw_paragraphs:
+            # 把段落內所有換行改成空格
+            single_line = para.replace('\r\n', ' ').replace('\r', ' ').replace('\n', ' ')
+            cleaned = single_line.strip()
+            if cleaned:
+                result.append(cleaned)
+        return result
+
     def _split_large_chunk(self, paragraphs: List[str], embeddings: List[List[float]]) -> (List[str], List[str]):
         """
         Splits 'paragraphs' by finding the least similar boundary using 'embeddings'
@@ -25,8 +41,9 @@ class Split:
             return paragraphs, []
 
         # Find the least similar consecutive boundary
-        split_index = self.utils.get_consecutive_least_similar(embeddings)
-
+        window_size = 3
+        split_index = self.utils.get_windowed_least_similar(embeddings, window_size=window_size)
+        
         left_part = paragraphs[:split_index + 1]
         right_part = paragraphs[split_index + 1:]
         return left_part, right_part
@@ -37,84 +54,76 @@ class Split:
         min_chunk_size: int = 300,
         max_chunk_size: int = 500
     ) -> List[str]:
-        """
-        Splits the input text into chunks of token-size between [min_chunk_size, max_chunk_size].
-        Once a chunk is in that range, we find the "least similar" boundary, store the left side,
-        and re-insert the right side for further splitting.
-        """
-        paragraphs = [p.strip() for p in extracted_text.split("\n\n") if p.strip()]
-        if not paragraphs:
+        
+        # 1) Build a deque of triples, so we never mutate three separate lists:
+        # paras = [p.strip() for p in extracted_text.split("\n\n") if p.strip()]
+        paras = self._custom_split(extracted_text)
+        
+        if not paras:
             return []
 
-        # Precompute once
-        paragraphs_tokens = [self.utils.count_tokens(p) for p in paragraphs]
-        paragraphs_embeddings = self.utils.list_embeddings(paragraphs)
-
-        final_chunks = []
-        idx = 0
-        n = len(paragraphs)
-
-        while idx < n:
-            chunk_paragraphs = []
-            chunk_embeddings = []
-            chunk_tokens = []  # Keep track of tokens in this chunk
-            current_tokens = 0
-
-            # 1) Accumulate until we at least exceed min_chunk_size or run out
-            while idx < n and current_tokens < min_chunk_size:
-                if current_tokens + paragraphs_tokens[idx] <= max_chunk_size:
-                    chunk_paragraphs.append(paragraphs[idx])
-                    chunk_embeddings.append(paragraphs_embeddings[idx])
-                    chunk_tokens.append(paragraphs_tokens[idx])
-                    current_tokens += paragraphs_tokens[idx]
-                    idx += 1
-                else:
-                    # This paragraph alone might exceed max_chunk_size -> handle as you see fit
+        tokens = [self.utils.count_tokens(p) for p in paras]
+        embs   = self.utils.list_embeddings(paras)
+        D: deque[Tuple[str,List[float],int]] = deque(
+            zip(paras, embs, tokens)
+        )
+
+        final_chunks: List[str] = []
+
+        # 2) As long as there’s anything left in D, build one chunk at a time:
+        while D:
+            cur_paras:    List[str]           = []
+            cur_embs:     List[List[float]]   = []
+            cur_tokens:   List[int]           = []
+            total_tokens = 0
+
+            # 2a) Guarantee we hit at least min_chunk_size
+            while D and total_tokens < min_chunk_size:
+                p, e, t = D.popleft()
+                # if even this one p would bust max, you might choose to take it alone
+                if total_tokens + t > max_chunk_size and total_tokens > 0:
+                    # push it back for the next round
+                    D.appendleft((p,e,t))
                     break
-
-            # If we haven't hit min_chunk_size but are out of paragraphs, store remainder and quit
-            if current_tokens < min_chunk_size and idx >= n:
-                if chunk_paragraphs:
-                    final_chunks.append(" ".join(chunk_paragraphs))
+                cur_paras.append(p)
+                cur_embs .append(e)
+                cur_tokens.append(t)
+                total_tokens += t
+
+            # if we ran out before min and have something -> emit it
+            if total_tokens < min_chunk_size and not D:
+                final_chunks.append(" ".join(cur_paras))
                 break
 
-            # 2) Keep adding while we're under max_chunk_size
-            while idx < n:
-                if current_tokens + paragraphs_tokens[idx] <= max_chunk_size:
-                    chunk_paragraphs.append(paragraphs[idx])
-                    chunk_embeddings.append(paragraphs_embeddings[idx])
-                    chunk_tokens.append(paragraphs_tokens[idx])
-                    current_tokens += paragraphs_tokens[idx]
-                    idx += 1
+            # 2b) Greedily fill until just under max_chunk_size
+            while D and total_tokens + D[0][2] <= max_chunk_size:
+                p, e, t = D.popleft()
+                cur_paras.append(p)
+                cur_embs .append(e)
+                cur_tokens.append(t)
+                total_tokens += t
+
+            # 3) Now we have between min and max tokens: split at the least-similar boundary
+            if cur_paras:
+                left, right = self._split_large_chunk(cur_paras, cur_embs)
+
+                # Count tokens in “left” to see if it meets min_chunk_size
+                left_token_count = sum(self.utils.count_tokens(p) for p in left)
+
+                if left_token_count >= min_chunk_size:
+                    # If left is big enough, emit it
+                    final_chunks.append(" ".join(left))
+
+                    # Push “right” (the remainder) back onto D for subsequent chunks
+                    for rp, re, rt in reversed(list(zip(
+                                    cur_paras[len(left):],
+                                    cur_embs  [len(left):],
+                                    cur_tokens[len(left):]
+                                ))):
+                        D.appendleft((rp, re, rt))
                 else:
-                    break
-
-            # Now we have between min_chunk_size and max_chunk_size tokens in 'chunk_paragraphs'
-            if chunk_paragraphs:
-                # 3) Split at the "least similar" boundary
-                left_part, right_part = self._split_large_chunk(
-                    chunk_paragraphs, chunk_embeddings
-                )
-
-                # We'll figure out how many paragraphs ended up in the left part
-                used_count = len(left_part)
-                leftover_count = len(right_part)
-
-                # Store left side
-                final_chunks.append(" ".join(left_part))
-
-                # If there's leftover, reinsert it into the main lists
-                if leftover_count > 0:
-                    # Slices for leftover
-                    leftover_embeddings = chunk_embeddings[used_count:]
-                    leftover_tokens = chunk_tokens[used_count:]
-
-                    # Re-insert them at index=idx
-                    paragraphs[idx:idx] = right_part
-                    paragraphs_embeddings[idx:idx] = leftover_embeddings
-                    paragraphs_tokens[idx:idx] = leftover_tokens
-
-                    # Recompute n, in case the paragraphs list has grown
-                    n = len(paragraphs)
+                    # If “left” is too small, just emit the entire cur_paras as one chunk
+                    final_chunks.append(" ".join(cur_paras))
+                    # (We do NOT push anything back, because cur_paras is fully consumed.)
 
-        return final_chunks
+        return final_chunks 

hie_rag/tree_index.py

@@ -10,8 +10,8 @@ from .utils import Utils
 
 
 class TreeIndex:
-    def __init__(self, base_url: str):
-        self.client = ChatOllama(model="llama3.2:latest")
+    def __init__(self, base_url: str, model="llama3.2:latest"):
+        self.client = ChatOllama(model=model)
         self.utils = Utils(base_url=base_url)
 
     def _convert_to_string(self, chunk_metadata: dict) -> str:
@@ -32,8 +32,8 @@ class TreeIndex:
 
             NOTE: 
             1. 請輸出繁體中文
-            2. The summary should be concise and capture the main points of the text.
-            3. The summary should be around 5-8 sentences long.
+            2. The summary should be concise with details and better than the individual summaries.
+            3. The summary should be long enough to cover all the main points of the text.
 
             Summaries:
             {summaries}

hie_rag/utils.py

@@ -12,7 +12,7 @@ from .ai_client import AiClient
 
 
 class Utils:
-    def __init__(self, base_url=None):
+    def __init__(self, base_url: str):
         # self.client = OpenAI(api_key=api_key)
         self.client = AiClient(base_url=base_url)
 
@@ -24,13 +24,16 @@ class Utils:
         # Accept both raw bytes and file-like objects with `.read()`
         if isinstance(uploaded_file, bytes):
             file_bytes = uploaded_file
+            suffix = ".bin"  # fallback generic extension
         elif hasattr(uploaded_file, "read"):
             file_bytes = uploaded_file.read()
+            filename = getattr(uploaded_file, "name", None) or getattr(uploaded_file, "filename", None)
+            suffix = os.path.splitext(filename)[-1] if filename else ".bin"
         else:
             raise TypeError("Unsupported file type: must be bytes or file-like object")
 
         # Write to temp file for MarkItDown to process
-        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
             temp_file_path = temp_file.name
             temp_file.write(file_bytes)
 
@@ -71,4 +74,49 @@ class Utils:
         # Find the index where consecutive similarity is lowest
         split_index = np.argmin(consecutive_similarities)
         
-        return split_index
+        return split_index
+    
+    def get_windowed_least_similar(
+        self,
+        embeddings: list,
+        window_size: int = 3
+    ) -> int:
+        """
+        對 embeddings 做滑動窗口：對每個可能的分割位置 i（0 <= i < len-1），
+        將 [max(0, i-window_size+1) .. i] 這 window_size 句平均後的向量
+        與 [i+1 .. min(len, i+window_size)] 這 window_size 句平均後的向量做 cosine 相似度，
+        回傳相似度最低的那個 i。
+        """
+        if len(embeddings) < 2:
+            # 根本沒得分割
+            return 0
+
+        # 把 list-of-lists 轉成 numpy array (shape: [n_sentences, dim_emb])
+        embs = np.array(embeddings)
+        n = embs.shape[0]
+
+        best_index = 0
+        lowest_sim = float('inf')
+
+        for i in range(n - 1):
+            # 前半段：從 pre_start 到 i (inclusive)
+            pre_start = max(0, i - window_size + 1)
+            pre_group = embs[pre_start : i + 1]   # shape: (<=window_size, dim)
+
+            # 後半段：從 i+1 到 post_end-1
+            post_end = min(n, i + 1 + window_size)
+            post_group = embs[i + 1 : post_end]   # shape: (<=window_size, dim)
+
+            # 計算平均向量
+            # （也可以改成加總：np.sum(...)；不過平均比較常見且 scale 感覺一致）
+            pre_avg = np.mean(pre_group, axis=0).reshape(1, -1)   # shape: (1, dim)
+            post_avg = np.mean(post_group, axis=0).reshape(1, -1) # shape: (1, dim)
+
+            # 計算 cosine similarity
+            sim = float(cosine_similarity(pre_avg, post_avg)[0][0])
+
+            if sim < lowest_sim:
+                lowest_sim = sim
+                best_index = i
+
+        return best_index

{hie_rag-0.2.0.dist-info → hie_rag-0.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hie_rag
-Version: 0.2.0
+Version: 0.2.2
 Summary: A hierarchical RAG framework for chunks retrieval.
 License: Apache License
         Version 2.0, January 2004

hie_rag-0.2.2.dist-info/RECORD

@@ -0,0 +1,14 @@
+hie_rag/__init__.py,sha256=p2glSTkCqGvMlcivcuKBStFh2C5adojaC9aGmF6nbhY,358
+hie_rag/ai_client.py,sha256=VbGQ0e3vZNn8W2YoR15Vvq2r-MUs-TBRNLGiImT4QxU,1000
+hie_rag/hie_rag.py,sha256=Nl_1WZM9IWhpNyZMvPzsae_u_xaCWEwrJgorZV-hp20,2741
+hie_rag/process.py,sha256=Z4qpNmxSsxUJgnqJtw8cYWJTS6SxhRR7F7eX_akyVCU,2427
+hie_rag/split.py,sha256=gEQVt57xWruT5e1psgSOnwuBrQngzri3S4H6ZvKzsw4,5082
+hie_rag/split_and_process.py,sha256=PkFlnOF7nW4Zs47JTsGF4AY9VDOXz1AtxG9Die8_mQk,572
+hie_rag/tree_index.py,sha256=iTa25ohMv5O0HYc5JtzIzVAIhNdVklYiAIJvqyE8sbM,2722
+hie_rag/utils.py,sha256=GwGiQj-zc8-U9UXOFHTKkjHWHx8YTYquR27gsXJgzCE,4687
+hie_rag/vectordb.py,sha256=iI73ujrONjDaHU66RNdHnD2PZWSppnjm0isIHPJEGAY,11068
+hie_rag-0.2.2.dist-info/licenses/LICENSE,sha256=IwAxruLb1UG8F0KZtfnV6MJq10FRAxWM-XOTWkWsJt4,632
+hie_rag-0.2.2.dist-info/METADATA,sha256=3vTI_zyvJxOOq8VrrchOAn0a7m8hwQPISnlholFi3u0,1698
+hie_rag-0.2.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+hie_rag-0.2.2.dist-info/top_level.txt,sha256=tN2S3VpMUl6oLWL9sN4xIh4o2na_zjnW8rHiwPFf0T8,8
+hie_rag-0.2.2.dist-info/RECORD,,

{hie_rag-0.2.0.dist-info → hie_rag-0.2.2.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.8.0)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

hie_rag-0.2.0.dist-info/RECORD

@@ -1,14 +0,0 @@
-hie_rag/__init__.py,sha256=p2glSTkCqGvMlcivcuKBStFh2C5adojaC9aGmF6nbhY,358
-hie_rag/ai_client.py,sha256=VbGQ0e3vZNn8W2YoR15Vvq2r-MUs-TBRNLGiImT4QxU,1000
-hie_rag/hie_rag.py,sha256=KB44QBz3tE0Eq_FJw9pvKynCfjyAuulaMFYKk6bzjug,2359
-hie_rag/process.py,sha256=D_vMnF84ingLb4_KoC77uLQXSa6FwEpR30RGukG2H9U,2414
-hie_rag/split.py,sha256=My7QQ_pPiJD0TvwRzm2MgonMMA79-r3Vifwp1xLWX4I,4905
-hie_rag/split_and_process.py,sha256=PkFlnOF7nW4Zs47JTsGF4AY9VDOXz1AtxG9Die8_mQk,572
-hie_rag/tree_index.py,sha256=TuRi9-M2aiD46ciS-iwIJYDc9nXq7i7mwxwVbMXk5Lo,2668
-hie_rag/utils.py,sha256=F5bqx147yT37z080MPWPrwzOa0tGEAWmvNFgjXpe4ZA,2729
-hie_rag/vectordb.py,sha256=iI73ujrONjDaHU66RNdHnD2PZWSppnjm0isIHPJEGAY,11068
-hie_rag-0.2.0.dist-info/licenses/LICENSE,sha256=IwAxruLb1UG8F0KZtfnV6MJq10FRAxWM-XOTWkWsJt4,632
-hie_rag-0.2.0.dist-info/METADATA,sha256=Oym7z46OyhT_Gp7unhX1rsYlFQi9UuOBU5VRsko1m_A,1698
-hie_rag-0.2.0.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
-hie_rag-0.2.0.dist-info/top_level.txt,sha256=tN2S3VpMUl6oLWL9sN4xIh4o2na_zjnW8rHiwPFf0T8,8
-hie_rag-0.2.0.dist-info/RECORD,,