hie-rag 0.2.0__tar.gz → 0.2.2__tar.gz

{hie_rag-0.2.0 → hie_rag-0.2.2}/PKG-INFO

@@ -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.0 → hie_rag-0.2.2}/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-0.2.0 → hie_rag-0.2.2}/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-0.2.2/hie_rag/split.py

@@ -0,0 +1,129 @@
+import re
+from collections import deque
+from typing import List, Tuple
+
+from .utils import Utils
+
+
+class Split:
+    def __init__(self, base_url: str = None):
+        """
+        Initializes the Split object with default or user-defined thresholds.
+        """
+        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'
+        (which are precomputed for these paragraphs only). Returns (left_part, right_part).
+        """
+        # If there are 0 or 1 paragraphs, no need to split
+        if len(paragraphs) < 2:
+            return paragraphs, []
+
+        # We'll assume 'embeddings' is already the same length as 'paragraphs'.
+        if len(embeddings) < 2:
+            # Can't compute consecutive similarities with fewer than 2 embeddings
+            return paragraphs, []
+
+        # Find the least similar consecutive boundary
+        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
+
+    def split(
+        self,
+        extracted_text: str,
+        min_chunk_size: int = 300,
+        max_chunk_size: int = 500
+    ) -> List[str]:
+        
+        # 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 []
+
+        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
+                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
+
+            # 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:
+                    # 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 

{hie_rag-0.2.0 → hie_rag-0.2.2}/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-0.2.0 → hie_rag-0.2.2}/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 → hie_rag-0.2.2}/hie_rag.egg-info/PKG-INFO

@@ -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.0 → hie_rag-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "hie_rag"
-version = "0.2.0"
+version = "0.2.2"
 description = "A hierarchical RAG framework for chunks retrieval."
 license = { file = "LICENSE" }
 dependencies = [

hie_rag-0.2.2/test/test-utils.py

@@ -0,0 +1,56 @@
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from dotenv import load_dotenv
+
+from hie_rag.utils import Utils
+
+load_dotenv()
+
+utils = Utils(base_url=os.getenv("BASE_URL"))
+
+with open("test2.pdf", "rb") as uploaded_file:
+    extracted_text = utils.extract_text(uploaded_file)
+
+# Count tokens for the first 100 words
+# result_count_tokens = utils.count_tokens(extracted_text[:100])
+result_count_tokens = utils.count_tokens(extracted_text)
+print(f"Token count: {result_count_tokens}")
+
+# # Get embeddings for the text slices
+# result_list_embeddings = utils.list_embeddings([
+#     extracted_text[:100],
+#     extracted_text[100:200],
+#     extracted_text[200:300]
+# ])
+
+# # Get the embedding for the first 100 words
+# result_get_embedding = utils.get_embedding(extracted_text[:100])
+
+# # Find the index of least similar consecutive embeddings
+# result_get_consecutive_least_similar = utils.get_consecutive_least_similar(result_list_embeddings)
+
+# # Write results to the text file
+# with open("test-utils-result-new", "w", encoding="utf-8") as file:
+#     file.write("Extracted Text:\n")
+#     file.write(extracted_text + "\n\n")
+#     file.write("====================================\n\n")
+    
+#     file.write("Count of Tokens (First 100 words):\n")
+#     file.write(str(result_count_tokens) + "\n\n")
+#     file.write("====================================\n\n")
+    
+#     file.write("List of Embeddings:\n")
+#     file.write(str(result_list_embeddings) + "\n\n")
+#     file.write("====================================\n\n")
+    
+#     file.write("Embedding of First 100 words:\n")
+#     file.write(str(result_get_embedding) + "\n\n")
+#     file.write("====================================\n\n")
+    
+#     file.write("Index of Least Similar Consecutive Embeddings:\n")
+#     file.write(str(result_get_consecutive_least_similar) + "\n")
+
+# print("Results written to a txt file.")

hie_rag-0.2.0/hie_rag/split.py

@@ -1,120 +0,0 @@
-from typing import List
-
-from .utils import Utils
-
-
-class Split:
-    def __init__(self, base_url: str = None):
-        """
-        Initializes the Split object with default or user-defined thresholds.
-        """
-        self.utils = Utils(base_url=base_url)
-
-    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'
-        (which are precomputed for these paragraphs only). Returns (left_part, right_part).
-        """
-        # If there are 0 or 1 paragraphs, no need to split
-        if len(paragraphs) < 2:
-            return paragraphs, []
-
-        # We'll assume 'embeddings' is already the same length as 'paragraphs'.
-        if len(embeddings) < 2:
-            # Can't compute consecutive similarities with fewer than 2 embeddings
-            return paragraphs, []
-
-        # Find the least similar consecutive boundary
-        split_index = self.utils.get_consecutive_least_similar(embeddings)
-
-        left_part = paragraphs[:split_index + 1]
-        right_part = paragraphs[split_index + 1:]
-        return left_part, right_part
-
-    def split(
-        self,
-        extracted_text: str,
-        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:
-            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
-                    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))
-                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
-                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)
-
-        return final_chunks

hie_rag-0.2.0/test/test-utils.py

@@ -1,54 +0,0 @@
-import os
-import sys
-
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
-
-from dotenv import load_dotenv
-
-from hie_rag.utils import Utils
-
-load_dotenv()
-
-utils = Utils(base_url=os.getenv("BASE_URL"))
-
-with open("test.pdf", "rb") as uploaded_file:
-    extracted_text = utils.extract_text(uploaded_file)
-
-# Count tokens for the first 100 words
-result_count_tokens = utils.count_tokens(extracted_text[:100])
-
-# Get embeddings for the text slices
-result_list_embeddings = utils.list_embeddings([
-    extracted_text[:100],
-    extracted_text[100:200],
-    extracted_text[200:300]
-])
-
-# Get the embedding for the first 100 words
-result_get_embedding = utils.get_embedding(extracted_text[:100])
-
-# Find the index of least similar consecutive embeddings
-result_get_consecutive_least_similar = utils.get_consecutive_least_similar(result_list_embeddings)
-
-# Write results to the text file
-with open("test-utils-result-new", "w", encoding="utf-8") as file:
-    file.write("Extracted Text:\n")
-    file.write(extracted_text + "\n\n")
-    file.write("====================================\n\n")
-    
-    file.write("Count of Tokens (First 100 words):\n")
-    file.write(str(result_count_tokens) + "\n\n")
-    file.write("====================================\n\n")
-    
-    file.write("List of Embeddings:\n")
-    file.write(str(result_list_embeddings) + "\n\n")
-    file.write("====================================\n\n")
-    
-    file.write("Embedding of First 100 words:\n")
-    file.write(str(result_get_embedding) + "\n\n")
-    file.write("====================================\n\n")
-    
-    file.write("Index of Least Similar Consecutive Embeddings:\n")
-    file.write(str(result_get_consecutive_least_similar) + "\n")
-
-print("Results written to a txt file.")