lionagi 0.0.312__py3-none-any.whl → 0.2.1__py3-none-any.whl

lionagi/libs/ln_image.py

@@ -0,0 +1,107 @@
+import base64
+import numpy as np
+from typing import Optional
+from .sys_util import SysUtil
+
+
+class ImageUtil:
+
+    @staticmethod
+    def preprocess_image(
+        image: np.ndarray, color_conversion_code: Optional[int] = None
+    ) -> np.ndarray:
+        SysUtil.check_import("cv2", pip_name="opencv-python")
+        import cv2
+
+        color_conversion_code = color_conversion_code or cv2.COLOR_BGR2RGB
+        return cv2.cvtColor(image, color_conversion_code)
+
+    @staticmethod
+    def encode_image_to_base64(image: np.ndarray, file_extension: str = ".jpg") -> str:
+        SysUtil.check_import("cv2", pip_name="opencv-python")
+        import cv2
+
+        success, buffer = cv2.imencode(file_extension, image)
+        if not success:
+            raise ValueError(f"Could not encode image to {file_extension} format.")
+        encoded_image = base64.b64encode(buffer).decode("utf-8")
+        return encoded_image
+
+    @staticmethod
+    def read_image_to_array(
+        image_path: str, color_flag: Optional[int] = None
+    ) -> np.ndarray:
+        SysUtil.check_import("cv2", pip_name="opencv-python")
+        import cv2
+
+        image = cv2.imread(image_path, color_flag)
+        color_flag = color_flag or cv2.IMREAD_COLOR
+        if image is None:
+            raise ValueError(f"Could not read image from path: {image_path}")
+        return image
+
+    @staticmethod
+    def read_image_to_base64(
+        image_path: str,
+        color_flag: Optional[int] = None,
+    ) -> str:
+        image_path = str(image_path)
+        image = ImageUtil.read_image_to_array(image_path, color_flag)
+
+        file_extension = "." + image_path.split(".")[-1]
+        return ImageUtil.encode_image_to_base64(image, file_extension)
+
+    # @staticmethod
+    # def encode_image(image_path):
+    #     with open(image_path, "rb") as image_file:
+    #         return base64.b64encode(image_file.read()).decode("utf-8")
+
+    @staticmethod
+    def calculate_image_token_usage_from_base64(image_base64: str, detail):
+        """
+        Calculate the token usage for processing OpenAI images from a base64-encoded string.
+
+        Parameters:
+        image_base64 (str): The base64-encoded string of the image.
+        detail (str): The detail level of the image, either 'low' or 'high'.
+
+        Returns:
+        int: The total token cost for processing the image.
+        """
+        import base64
+        from io import BytesIO
+        from PIL import Image
+
+        # Decode the base64 string to get image data
+        if "data:image/jpeg;base64," in image_base64:
+            image_base64 = image_base64.split("data:image/jpeg;base64,")[1]
+            image_base64.strip("{}")
+
+        image_data = base64.b64decode(image_base64)
+        image = Image.open(BytesIO(image_data))
+
+        # Get image dimensions
+        width, height = image.size
+
+        if detail == "low":
+            return 85
+
+        # Scale to fit within a 2048 x 2048 square
+        max_dimension = 2048
+        if width > max_dimension or height > max_dimension:
+            scale_factor = max_dimension / max(width, height)
+            width = int(width * scale_factor)
+            height = int(height * scale_factor)
+
+        # Scale such that the shortest side is 768px
+        min_side = 768
+        if min(width, height) > min_side:
+            scale_factor = min_side / min(width, height)
+            width = int(width * scale_factor)
+            height = int(height * scale_factor)
+
+        # Calculate the number of 512px squares
+        num_squares = (width // 512) * (height // 512)
+        token_cost = 170 * num_squares + 85
+
+        return token_cost

lionagi/libs/ln_knowledge_graph.py

@@ -0,0 +1,405 @@
+import math
+from lionagi.libs import CallDecorator as cd
+
+
+class KnowledgeBase:
+    """
+    A class to represent a Knowledge Base (KB) containing entities, relations, and sources.
+
+    Attributes:
+        entities (dict): A dictionary of entities in the KB, where the keys are entity titles, and the values are
+                         entity information (excluding the title).
+        relations (list): A list of relations in the KB, where each relation is a dictionary containing information
+                          about the relation (head, type, tail) and metadata (article_url and spans).
+        sources (dict): A dictionary of information about the sources of relations, where the keys are article URLs,
+                       and the values are source data (article_title and article_publish_date).
+
+    Methods:
+        merge_with_kb(kb2): Merge another Knowledge Base (kb2) into this KB.
+        are_relations_equal(r1, r2): Check if two relations (r1 and r2) are equal.
+        exists_relation(r1): Check if a relation (r1) already exists in the KB.
+        merge_relations(r2): Merge the information from relation r2 into an existing relation in the KB.
+        get_wikipedia_data(candidate_entity): Get data for a candidate entity from Wikipedia.
+        add_entity(e): Add an entity to the KB.
+        add_relation(r, article_title, article_publish_date): Add a relation to the KB.
+        print(): Print the entities, relations, and sources in the KB.
+        extract_relations_from_model_output(text): Extract relations from the model output text.
+
+    """
+
+    def __init__(self):
+        """
+        Initialize an empty Knowledge Base (KB) with empty dictionaries for entities, relations, and sources.
+        """
+        self.entities = {}  # { entity_title: {...} }
+        self.relations = []  # [ head: entity_title, type: ..., tail: entity_title,
+        # meta: { article_url: { spans: [...] } } ]
+        self.sources = {}  # { article_url: {...} }
+
+    def merge_with_kb(self, kb2):
+        """
+        Merge another Knowledge Base (KB) into this KB.
+
+        Args:
+            kb2 (KnowledgeBase): The Knowledge Base (KB) to merge into this KB.
+        """
+        for r in kb2.relations:
+            article_url = list(r["meta"].keys())[0]
+            source_data = kb2.sources[article_url]
+            self.add_relation(
+                r, source_data["article_title"], source_data["article_publish_date"]
+            )
+
+    def are_relations_equal(self, r1, r2):
+        """
+        Check if two relations (r1 and r2) are equal.
+
+        Args:
+            r1 (dict): The first relation to compare.
+            r2 (dict): The second relation to compare.
+
+        Returns:
+            bool: True if the relations are equal, False otherwise.
+        """
+        return all(r1[attr] == r2[attr] for attr in ["head", "type", "tail"])
+
+    def exists_relation(self, r1):
+        """
+        Check if a relation (r1) already exists in the KB.
+
+        Args:
+            r1 (dict): The relation to check for existence in the KB.
+
+        Returns:
+            bool: True if the relation exists in the KB, False otherwise.
+        """
+        return any(self.are_relations_equal(r1, r2) for r2 in self.relations)
+
+    def merge_relations(self, r2):
+        """
+        Merge the information from relation r2 into an existing relation in the KB.
+
+        Args:
+            r2 (dict): The relation to merge into an existing relation in the KB.
+        """
+        r1 = [r for r in self.relations if self.are_relations_equal(r2, r)][0]
+
+        # if different article
+        article_url = list(r2["meta"].keys())[0]
+        if article_url not in r1["meta"]:
+            r1["meta"][article_url] = r2["meta"][article_url]
+
+        # if existing article
+        else:
+            spans_to_add = [
+                span
+                for span in r2["meta"][article_url]["spans"]
+                if span not in r1["meta"][article_url]["spans"]
+            ]
+            r1["meta"][article_url]["spans"] += spans_to_add
+
+    @cd.cache(maxsize=10000)
+    def get_wikipedia_data(self, candidate_entity):
+        """
+        Get data for a candidate entity from Wikipedia.
+
+        Args:
+            candidate_entity (str): The candidate entity title.
+
+        Returns:
+            dict: A dictionary containing information about the candidate entity (title, url, summary).
+                  None if the entity does not exist in Wikipedia.
+        """
+        try:
+            from lionagi.libs import SysUtil
+
+            SysUtil.check_import("wikipedia")
+            import wikipedia  # type: ignore
+        except Exception as e:
+            raise Exception("wikipedia package is not installed {e}")
+
+        try:
+            page = wikipedia.page(candidate_entity, auto_suggest=False)
+            entity_data = {
+                "title": page.title,
+                "url": page.url,
+                "summary": page.summary,
+            }
+            return entity_data
+        except:
+            return None
+
+    def add_entity(self, e):
+        """
+        Add an entity to the KB.
+
+        Args:
+            e (dict): A dictionary containing information about the entity (title and additional attributes).
+        """
+        self.entities[e["title"]] = {k: v for k, v in e.items() if k != "title"}
+
+    def add_relation(self, r, article_title, article_publish_date):
+        """
+        Add a relation to the KB.
+
+        Args:
+            r (dict): A dictionary containing information about the relation (head, type, tail, and metadata).
+            article_title (str): The title of the article containing the relation.
+            article_publish_date (str): The publish date of the article.
+        """
+        # check on wikipedia
+        candidate_entities = [r["head"], r["tail"]]
+        entities = [self.get_wikipedia_data(ent) for ent in candidate_entities]
+
+        # if one entity does not exist, stop
+        if any(ent is None for ent in entities):
+            return
+
+        # manage new entities
+        for e in entities:
+            self.add_entity(e)
+
+        # rename relation entities with their wikipedia titles
+        r["head"] = entities[0]["title"]
+        r["tail"] = entities[1]["title"]
+
+        # add source if not in kb
+        article_url = list(r["meta"].keys())[0]
+        if article_url not in self.sources:
+            self.sources[article_url] = {
+                "article_title": article_title,
+                "article_publish_date": article_publish_date,
+            }
+
+        # manage new relation
+        if not self.exists_relation(r):
+            self.relations.append(r)
+        else:
+            self.merge_relations(r)
+
+    def print(self):
+        """
+        Print the entities, relations, and sources in the KB.
+
+        Returns:
+            None
+        """
+        print("Entities:")
+        for e in self.entities.items():
+            print(f"  {e}")
+        print("Relations:")
+        for r in self.relations:
+            print(f"  {r}")
+        print("Sources:")
+        for s in self.sources.items():
+            print(f"  {s}")
+
+    @staticmethod
+    def extract_relations_from_model_output(text):
+        """
+        Extract relations from the model output text.
+
+        Args:
+            text (str): The model output text containing relations.
+
+        Returns:
+            list: A list of dictionaries, where each dictionary represents a relation (head, type, tail).
+        """
+        relations = []
+        relation, subject, relation, object_ = "", "", "", ""
+        text = text.strip()
+        current = "x"
+        text_replaced = text.replace("<s>", "").replace("<pad>", "").replace("</s>", "")
+        for token in text_replaced.split():
+            if token == "<triplet>":
+                current = "t"
+                if relation != "":
+                    relations.append(
+                        {
+                            "head": subject.strip(),
+                            "type": relation.strip(),
+                            "tail": object_.strip(),
+                        }
+                    )
+                    relation = ""
+                subject = ""
+            elif token == "<subj>":
+                current = "s"
+                if relation != "":
+                    relations.append(
+                        {
+                            "head": subject.strip(),
+                            "type": relation.strip(),
+                            "tail": object_.strip(),
+                        }
+                    )
+                object_ = ""
+            elif token == "<obj>":
+                current = "o"
+                relation = ""
+            else:
+                if current == "t":
+                    subject += " " + token
+                elif current == "s":
+                    object_ += " " + token
+                elif current == "o":
+                    relation += " " + token
+        if subject != "" and relation != "" and object_ != "":
+            relations.append(
+                {
+                    "head": subject.strip(),
+                    "type": relation.strip(),
+                    "tail": object_.strip(),
+                }
+            )
+        return relations
+
+
+class KGTripletExtractor:
+    """
+    A class to perform knowledge graph triplet extraction from text using a pre-trained model.
+
+    Methods:
+        text_to_wiki_kb(text, model=None, tokenizer=None, device='cpu', span_length=512,
+                        article_title=None, article_publish_date=None, verbose=False):
+            Extract knowledge graph triplets from text and create a KnowledgeBase (KB) containing entities and relations.
+
+    """
+
+    @staticmethod
+    def text_to_wiki_kb(
+        text,
+        model=None,
+        tokenizer=None,
+        device="cpu",
+        span_length=512,
+        article_title=None,
+        article_publish_date=None,
+        verbose=False,
+    ):
+        from lionagi.integrations.bridge.transformers_.install_ import (
+            install_transformers,
+        )
+
+        try:
+            from transformers import AutoModelForSeq2SeqLM, AutoTokenizer  # type: ignore
+        except ImportError:
+            install_transformers()
+        from transformers import AutoModelForSeq2SeqLM, AutoTokenizer  # type: ignore
+        import torch  # type: ignore
+
+        """
+        Extract knowledge graph triplets from text and create a KnowledgeBase (KB) containing entities and relations.
+
+        Args:
+            text (str): The input text from which triplets will be extracted.
+            model (AutoModelForSeq2SeqLM, optional): The pre-trained model for triplet extraction. Defaults to None.
+            tokenizer (AutoTokenizer, optional): The tokenizer for the model. Defaults to None.
+            device (str, optional): The device to run the model on (e.g., 'cpu', 'cuda'). Defaults to 'cpu'.
+            span_length (int, optional): The maximum span length for input text segmentation. Defaults to 512.
+            article_title (str, optional): The title of the article containing the input text. Defaults to None.
+            article_publish_date (str, optional): The publish date of the article. Defaults to None.
+            verbose (bool, optional): Whether to enable verbose mode for debugging. Defaults to False.
+
+        Returns:
+            KnowledgeBase: A KnowledgeBase (KB) containing extracted entities, relations, and sources.
+
+        """
+
+        if not any([model, tokenizer]):
+            tokenizer = AutoTokenizer.from_pretrained("Babelscape/rebel-large")
+            model = AutoModelForSeq2SeqLM.from_pretrained("Babelscape/rebel-large")
+            model.to(device)
+
+        inputs = tokenizer([text], return_tensors="pt")
+
+        num_tokens = len(inputs["input_ids"][0])
+        if verbose:
+            print(f"Input has {num_tokens} tokens")
+        num_spans = math.ceil(num_tokens / span_length)
+        if verbose:
+            print(f"Input has {num_spans} spans")
+        overlap = math.ceil(
+            (num_spans * span_length - num_tokens) / max(num_spans - 1, 1)
+        )
+        spans_boundaries = []
+        start = 0
+        for i in range(num_spans):
+            spans_boundaries.append(
+                [start + span_length * i, start + span_length * (i + 1)]
+            )
+            start -= overlap
+        if verbose:
+            print(f"Span boundaries are {spans_boundaries}")
+
+        # transform input with spans
+        tensor_ids = [
+            inputs["input_ids"][0][boundary[0] : boundary[1]]
+            for boundary in spans_boundaries
+        ]
+        tensor_masks = [
+            inputs["attention_mask"][0][boundary[0] : boundary[1]]
+            for boundary in spans_boundaries
+        ]
+
+        inputs = {
+            "input_ids": torch.stack(tensor_ids).to(device),
+            "attention_mask": torch.stack(tensor_masks).to(device),
+        }
+
+        # generate relations
+        num_return_sequences = 3
+        gen_kwargs = {
+            "max_length": 512,
+            "length_penalty": 0,
+            "num_beams": 3,
+            "num_return_sequences": num_return_sequences,
+        }
+        generated_tokens = model.generate(
+            **inputs,
+            **gen_kwargs,
+        )
+
+        # decode relations
+        decoded_preds = tokenizer.batch_decode(
+            generated_tokens, skip_special_tokens=False
+        )
+
+        # create kb
+        kb = KnowledgeBase()
+        i = 0
+        for sentence_pred in decoded_preds:
+            current_span_index = i // num_return_sequences
+            relations = KnowledgeBase.extract_relations_from_model_output(sentence_pred)
+            for relation in relations:
+                relation["meta"] = {
+                    "article_url": {"spans": [spans_boundaries[current_span_index]]}
+                }
+                kb.add_relation(relation, article_title, article_publish_date)
+            i += 1
+        return kb
+
+
+class KGraph:
+    """
+    A class representing a Knowledge Graph (KGraph) for extracting relations from text.
+
+    Methods:
+        text_to_wiki_kb(text, model=None, tokenizer=None, device='cpu', span_length=512, article_title=None,
+                        article_publish_date=None, verbose=False):
+            Extract relations from input text and create a Knowledge Base (KB) containing entities and relations.
+    """
+
+    @staticmethod
+    def text_to_wiki_kb(text, **kwargs):
+        """
+        Extract relations from input text and create a Knowledge Base (KB) containing entities and relations.
+
+        Args:
+            text (str): The input text from which relations are extracted.
+            **kwargs: Additional keyword arguments passed to the underlying extraction method.
+
+        Returns:
+            KnowledgeBase: A Knowledge Base (KB) containing entities and relations extracted from the input text.
+        """
+        return KGTripletExtractor.text_to_wiki_kb(text, **kwargs)

lionagi/libs/ln_nested.py

@@ -1,3 +1,19 @@
+"""
+Copyright 2024 HaiyangLi
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
 from collections import defaultdict
 from itertools import chain
 from typing import Any, Generator, Callable
@@ -52,7 +68,7 @@ def nset(nested_structure: dict | list, indices: list[int | str], value: Any) ->
 def nget(
     nested_structure: dict | list,
     indices: list[int | str],
-    default: Any | None = None,
+    default=...,
 ) -> Any:
     """
     retrieves a value from a nested list or dictionary structure, with an option to
@@ -98,12 +114,12 @@ def nget(
             return target_container[last_index]
         elif isinstance(target_container, dict) and last_index in target_container:
             return target_container[last_index]
-        elif default is not None:
+        elif default is not ...:
             return default
         else:
             raise LookupError("Target not found and no default value provided.")
     except (IndexError, KeyError, TypeError):
-        if default is not None:
+        if default is not ...:
             return default
         else:
             raise LookupError("Target not found and no default value provided.")
@@ -116,7 +132,7 @@ def nmerge(
     *,
     overwrite: bool = False,
     dict_sequence: bool = False,
-    sequence_separator: str = "_",
+    sequence_separator: str = "[^_^]",
     sort_list: bool = False,
     custom_sort: Callable[[Any], Any] | None = None,
 ) -> dict | list:
@@ -176,7 +192,7 @@ def flatten(
     /,
     *,
     parent_key: str = "",
-    sep: str = "_",
+    sep: str = "[^_^]",
     max_depth: int | None = None,
     inplace: bool = False,
     dict_only: bool = False,
@@ -238,7 +254,7 @@ def unflatten(
     flat_dict: dict[str, Any],
     /,
     *,
-    sep: str = "_",
+    sep: str = "[^_^]",
     custom_logic: Callable[[str], Any] | None = None,
     max_depth: int | None = None,
 ) -> dict | list:
@@ -330,7 +346,7 @@ def ninsert(
     indices: list[str | int],
     value: Any,
     *,
-    sep: str = "_",
+    sep: str = "[^_^]",
     max_depth: int | None = None,
     current_depth: int = 0,
 ) -> None:
@@ -393,12 +409,11 @@ def ninsert(
         nested_structure[last_part] = value
 
 
-# noinspection PyDecorator
 def get_flattened_keys(
     nested_structure: Any,
     /,
     *,
-    sep: str = "_",
+    sep: str = "[^_^]",
     max_depth: int | None = None,
     dict_only: bool = False,
     inplace: bool = False,
@@ -448,7 +463,7 @@ def _dynamic_flatten_in_place(
     /,
     *,
     parent_key: str = "",
-    sep: str = "_",
+    sep: str = "[^_^]",
     max_depth: int | None = None,
     current_depth: int = 0,
     dict_only: bool = False,
@@ -581,7 +596,7 @@ def _deep_update(original: dict, update: dict) -> dict:
 def _dynamic_flatten_generator(
     nested_structure: Any,
     parent_key: tuple[str, ...],
-    sep: str = "_",
+    sep: str = "[^_^]",
     max_depth: int | None = None,
     current_depth: int = 0,
     dict_only: bool = False,