hie-rag 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

hie_rag/__init__.py

@@ -1,5 +1,4 @@
 # __init__.py
-from .generate import Generate
 from .hie_rag import HieRag
 from .process import Process
 from .split import Split
@@ -7,3 +6,13 @@ from .split_and_process import SplitAndProcess
 from .tree_index import TreeIndex
 from .utils import Utils
 from .vectordb import Vectordb
+
+__all__ = [
+    "HieRag",
+    "Process",
+    "Split",
+    "SplitAndProcess",
+    "TreeIndex",
+    "Utils",
+    "Vectordb",
+]

hie_rag/hie_rag.py

@@ -38,7 +38,6 @@ class HieRag:
             "chunk_count": len(tree_index.get("chunks", [])),
         }
 
-
     def get_summary(self, file_id):
         return self.vector_db.get_summary(file_id)
 
@@ -55,35 +54,4 @@ class HieRag:
         return self.vector_db.query_summaries_by_text(query_text, n_results=n_results)
     
     def query_chunks_by_text(self, query_text: str, file_id: str, n_results=5):
-        return self.vector_db.query_chunks_by_text(query_text, file_id=file_id, n_results=n_results)
-    
-    # def query(self, query_text: str, n_results=5):
-    #     """
-    #     This n_result is for the chunks.
-    #     """
-    #     print("The summary is querying...")
-    #     query_summary_result = self.vector_db.query_summaries_by_text(query_text)
-    #     if not query_summary_result["metadatas"]:
-    #         return "No results found"
-
-    #     file_id = query_summary_result["metadatas"][0][0]["file_id"]
-    #     summary = query_summary_result["metadatas"][0][0]["summary"]
-    #     keywords = query_summary_result["metadatas"][0][0]["keywords"]
-
-    #     print("The chunks are querying...")
-    #     query_chunks_result = self.vector_db.query_chunks_by_text(query_text, file_id=file_id, n_results=n_results)
-
-    #     if not query_chunks_result["metadatas"]:
-    #         return "No results found"
-        
-    #     chunks = query_chunks_result["metadatas"][0]
-        
-    #     data = {
-    #         "file_id": file_id,
-    #         "summary": summary,
-    #         "keywords": keywords,
-    #         "chunks": chunks
-    #     }
-
-    #     return data
-
+        return self.vector_db.query_chunks_by_text(query_text, file_id=file_id, n_results=n_results)

{hie_rag-0.1.2.dist-info → hie_rag-0.1.3.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: hie_rag
-Version: 0.1.2
-Summary: A simple hierarchical RAG model
+Version: 0.1.3
+Summary: A hierarchical RAG framework for chunks retrieval.
 License: Apache License
         Version 2.0, January 2004
         http://www.apache.org/licenses/

{hie_rag-0.1.2.dist-info → hie_rag-0.1.3.dist-info}/RECORD

@@ -1,15 +1,14 @@
-hie_rag/__init__.py,sha256=mPaUHUxcZlu5RtGcY0NqZmUaeogsG-9vMboTCEP_yoU,264
+hie_rag/__init__.py,sha256=p2glSTkCqGvMlcivcuKBStFh2C5adojaC9aGmF6nbhY,358
 hie_rag/app.py,sha256=jZkGEIXhYL2mY3KhixXFqvkOn8r0Cdav3EZxlChvKDA,2636
-hie_rag/generate.py,sha256=qNRiRQMUWPZenhPRvtBzkrA7LBkXIQXnHlA0ICiunI4,2656
-hie_rag/hie_rag.py,sha256=rpQfbcPVYaeA2RCEcOxlwovYPsBYRzFYlfO9WP4piSo,3442
+hie_rag/hie_rag.py,sha256=h5EcGcxbcGm6-jB3lr_EIuZ-0wEQFJTF1xffzQKDJUI,2353
 hie_rag/process.py,sha256=JaL8i1IZckeeaHsNSYiUIlYRsRRB73E9QqLCSh09JHA,2434
 hie_rag/split.py,sha256=st_bZ4UaKUOXbxUIDobfG1IsW5vC9rHeyo4LXprfKrk,4904
 hie_rag/split_and_process.py,sha256=eRMiBYBZWUo3ljFasZGAOSP_6_adiwBD094DZJfVQDk,565
 hie_rag/tree_index.py,sha256=5rCoCCO14KLFvRzeOGB08mAnd6d3p7dl4h4jGQqF13A,2688
 hie_rag/utils.py,sha256=cxYLNch5CVgnpuD3ScVoJMP8Kp0_Ni3grF5tV1_sCOM,2769
 hie_rag/vectordb.py,sha256=UVdAinxUDhDqwbFbeXaLVdzN6uC4nu5l7rWi600d8BU,11065
-hie_rag-0.1.2.dist-info/licenses/LICENSE,sha256=IwAxruLb1UG8F0KZtfnV6MJq10FRAxWM-XOTWkWsJt4,632
-hie_rag-0.1.2.dist-info/METADATA,sha256=utsAOnIm1VW1r66Va8-plF5jWQVPDDWmES_PTwac3No,1680
-hie_rag-0.1.2.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-hie_rag-0.1.2.dist-info/top_level.txt,sha256=tN2S3VpMUl6oLWL9sN4xIh4o2na_zjnW8rHiwPFf0T8,8
-hie_rag-0.1.2.dist-info/RECORD,,
+hie_rag-0.1.3.dist-info/licenses/LICENSE,sha256=IwAxruLb1UG8F0KZtfnV6MJq10FRAxWM-XOTWkWsJt4,632
+hie_rag-0.1.3.dist-info/METADATA,sha256=U_s4BPalfUt8xQWqj1mHNJJC7IEPZuXSmeVaHOBNhn4,1699
+hie_rag-0.1.3.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+hie_rag-0.1.3.dist-info/top_level.txt,sha256=tN2S3VpMUl6oLWL9sN4xIh4o2na_zjnW8rHiwPFf0T8,8
+hie_rag-0.1.3.dist-info/RECORD,,

hie_rag/generate.py

@@ -1,61 +0,0 @@
-from typing import Dict, List
-
-from langchain_core.prompts import PromptTemplate
-from langchain_openai import ChatOpenAI
-from pydantic import Field
-from typing_extensions import TypedDict
-
-
-class Generate:
-    def __init__(self, api_key: str):   
-        self.client = ChatOpenAI(temperature=0, model="gpt-4o", api_key=api_key)
-
-    def generate(self, content: str, possible_reference: str) -> Dict:
-        """Generate data for finetuning"""
-        prompt = PromptTemplate(
-            template="""
-            你是一個資料生成器，負責生成用於微調模型的資料集。
-            你的工作是閱讀以下內容，並生成一系列人類可能給出的指令（instruction）以及對應的詳細回應（response）。
-            instruction 可能會是一個問題、請求整理或者整理內容等指令。
-
-            注意事項：
-            1. 請輸出繁體中文。
-            2. 請務必只生成與內容相關的指令與回應。
-            3. 如果不確定內容在講什麼，可以參考「可能參考資料（Possible Reference）」來幫助理解。
-            4.「可能參考資料」只是可能幫助你理解的參考來源。
-            5. 不要捏造答案，如果真的不知道，就不要亂寫。
-
-            Content:
-            {content}
-
-            Possible Reference:
-            {possible_reference}
-
-            """,
-            input_variables=["content", "possible_reference"],
-        )
-        class InstructionResponse(TypedDict):
-            instruction: str = Field(
-                description="An instruction that a human might provide based on the content.",
-            )
-            response: str = Field(
-                description="The corresponding response to the instruction.",
-            )
-            used_reference: bool = Field(
-                description="Indicates whether the possible reference was used to generate this pair. True if the Possible Reference is relavent and useful, False otherwise.",
-            )
-            reference_usage: str = Field(
-                description="Explanation of how the reference was used, if it was used.",
-            )
-
-        class Dataset(TypedDict):
-            dataset: List[InstructionResponse]
-            content_analysis: str = Field(
-                description="Brief analysis of whether and how the reference helped with understanding the content.",
-            )
-
-        model = self.client
-        llm_with_tool = model.with_structured_output(Dataset)
-        chain = prompt | llm_with_tool
-        
-        return chain.invoke({"content": content, "possible_reference": possible_reference})