llm-ie 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

llm_ie/data_types.py

@@ -1,4 +1,5 @@
-from typing import List, Dict, Iterable
+from typing import List, Dict, Tuple, Iterable, Callable
+import importlib.util
 import json
 
 
@@ -281,3 +282,82 @@ class LLMInformationExtractionDocument:
                         json_file, indent=4)
             json_file.flush()
             
+
+    def _viz_preprocess(self) -> Tuple:
+        """
+        This method preprocesses the entities and relations for visualization.
+        """
+        if importlib.util.find_spec("ie_viz") is None:
+            raise ImportError("ie_viz not found. Please install ie_viz (```pip install ie-viz```).")
+
+        if self.has_frame():
+            entities = [{"entity_id": frame.frame_id, "start": frame.start, "end": frame.end, "attr": frame.attr} for frame in self.frames]
+        else:
+            raise ValueError("No frames in the document.")
+        
+        if self.has_relation():
+            relations = []
+            for relation in self.relations:
+                rel = {"entity_1_id": relation['frame_1'], "entity_2_id": relation['frame_2']}
+                relations.append(rel)        
+        else:
+            relations = None
+
+        return entities, relations
+
+
+    def viz_serve(self, host: str = '0.0.0.0', port: int = 5000, theme:str = "light", 
+                  color_attr_key:str=None, color_map_func:Callable=None):
+        """
+        This method serves a visualization App of the document.
+
+        Parameters:
+        -----------
+        host : str, Optional
+            The host address to run the server on.
+        port : int, Optional
+            The port number to run the server on.
+        theme : str, Optional
+            The theme of the visualization. Must be either "light" or "dark".
+        color_attr_key : str, Optional
+            The attribute key to be used for coloring the entities.
+        color_map_func : Callable, Optional
+            The function to be used for mapping the entity attributes to colors. When provided, the color_attr_key and 
+            theme will be overwritten. The function must take an entity dictionary as input and return a color string (hex).
+        """
+        entities, relations = self._viz_preprocess()
+        from ie_viz import serve
+
+        serve(text=self.text,
+              entities=entities,
+              relations=relations,
+              host=host,
+              port=port,
+              theme=theme,
+              color_attr_key=color_attr_key,
+              color_map_func=color_map_func)
+        
+    
+    def viz_render(self, theme:str = "light", color_attr_key:str=None, color_map_func:Callable=None) -> str:
+        """
+        This method renders visualization html of the document.
+
+        Parameters:
+        -----------
+        theme : str, Optional
+            The theme of the visualization. Must be either "light" or "dark".
+        color_attr_key : str, Optional
+            The attribute key to be used for coloring the entities.
+        color_map_func : Callable, Optional
+            The function to be used for mapping the entity attributes to colors. When provided, the color_attr_key and 
+            theme will be overwritten. The function must take an entity dictionary as input and return a color string (hex).
+        """
+        entities, relations = self._viz_preprocess()
+        from ie_viz import render
+
+        return render(text=self.text,
+                      entities=entities,
+                      relations=relations,
+                      theme=theme,
+                      color_attr_key=color_attr_key,
+                      color_map_func=color_map_func)

llm_ie/extractors.py

@@ -896,7 +896,7 @@ class MultiClassRelationExtractor(RelationExtractor):
         if len(rel_json) > 0:
             if "RelationType" in rel_json[0]:
                 rel = rel_json[0]["RelationType"]
-                if rel in pos_rel_types:
+                if rel in pos_rel_types or rel == "No Relation":
                     return rel_json[0]["RelationType"]
                 else:
                     warnings.warn(f'Extracted relation type "{rel}", which is not in the return of possible_relation_types_func: {pos_rel_types}.'+ \

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: llm-ie
-Version: 0.2.0
+Version: 0.2.2
 Summary: An LLM-powered tool that transforms everyday language into robust information extraction pipelines.
 License: MIT
 Author: Enshuo (David) Hsu

{llm_ie-0.2.0.dist-info → llm_ie-0.2.2.dist-info}/RECORD

@@ -6,10 +6,10 @@ llm_ie/asset/prompt_guide/BinaryRelationExtractor_prompt_guide.txt,sha256=z9Xg0f
 llm_ie/asset/prompt_guide/MultiClassRelationExtractor_prompt_guide.txt,sha256=D5DphUHw8SUERUVdcIjUynuTmYJa6-PwBlF7FzxNsvQ,2276
 llm_ie/asset/prompt_guide/ReviewFrameExtractor_prompt_guide.txt,sha256=XbnU8byLGGUA3A3lT0bb2Hw-ggzhcqD3ZuKzduod2ww,1944
 llm_ie/asset/prompt_guide/SentenceFrameExtractor_prompt_guide.txt,sha256=8nj9OLPJMtr9Soi5JU3Xk-HC7pKNoI54xA_A4u7I5j4,2620
-llm_ie/data_types.py,sha256=2RKP4wXDuku-Tn4s8uzzUFavG1fZ2e47SaY8oL57LsI,10923
+llm_ie/data_types.py,sha256=hPz3WOeAzfn2QKmb0CxHmRdQWZQ4G9zq8U-RJBVFdYk,14329
 llm_ie/engines.py,sha256=m9ytGUX61jEy9SmVHbb90mrfGMAwC6dV-v7Jke1U7Ho,9296
-llm_ie/extractors.py,sha256=i0m8uFaKXiVY1ucjvzbUFbV1slPYfZ3EGOZrolnFVHA,44079
+llm_ie/extractors.py,sha256=EVuHqW1lW0RpGmnRNmX4ih6ppfvy2gYOAOgc8Pngfkw,44103
 llm_ie/prompt_editor.py,sha256=dbu7A3O7O7Iw2v-xCgrTFH1-wTLAGf4SHDqdeS-He2Q,1869
-llm_ie-0.2.0.dist-info/METADATA,sha256=9CPC3OAd2J0nROZ7z8DI7lvGKOO2H2uAnVQv_YDFItg,40052
-llm_ie-0.2.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-llm_ie-0.2.0.dist-info/RECORD,,
+llm_ie-0.2.2.dist-info/METADATA,sha256=V7gi1wZN_FlaW2ZfdsGKtEGi2s-AjIvssrkjjd7ubWU,40052
+llm_ie-0.2.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+llm_ie-0.2.2.dist-info/RECORD,,