llm-ie 0.2.0__tar.gz → 0.2.2__tar.gz

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

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

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "llm-ie"
-version = "0.2.0"
+version = "0.2.2"
 description = "An LLM-powered tool that transforms everyday language into robust information extraction pipelines."
 authors = ["Enshuo (David) Hsu"]
 license = "MIT"

{llm_ie-0.2.0 → llm_ie-0.2.2}/src/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-0.2.0 → llm_ie-0.2.2}/src/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}.'+ \