lionagi 0.2.11__py3-none-any.whl → 0.3.1__py3-none-any.whl

lionagi/integrations/storage/to_csv.py

@@ -3,7 +3,10 @@ from pathlib import Path
 
 import pandas as pd
 
-from lionagi.integrations.storage.storage_util import output_edge_list, output_node_list
+from lionagi.integrations.storage.storage_util import (
+    output_edge_list,
+    output_node_list,
+)
 
 
 def _output_csv(
@@ -27,7 +30,10 @@ def _output_csv(
     Returns:
         None: This function does not return a value but outputs a ZIP file containing the CSVs.
     """
-    tables = {"Nodes": pd.DataFrame(node_list), "Edges": pd.DataFrame(edge_list)}
+    tables = {
+        "Nodes": pd.DataFrame(node_list),
+        "Edges": pd.DataFrame(edge_list),
+    }
     if edge_cls_list:
         tables["EdgesCondClass"] = pd.DataFrame(edge_cls_list)
     for i in node_dict:

lionagi/integrations/storage/to_excel.py

@@ -1,6 +1,9 @@
 import pandas as pd
 
-from lionagi.integrations.storage.storage_util import output_edge_list, output_node_list
+from lionagi.integrations.storage.storage_util import (
+    output_edge_list,
+    output_node_list,
+)
 from lionagi.libs import SysUtil
 
 
@@ -38,7 +41,10 @@ def _output_excel(
 
     structure_id = ""
 
-    tables = {"Nodes": pd.DataFrame(node_list), "Edges": pd.DataFrame(edge_list)}
+    tables = {
+        "Nodes": pd.DataFrame(node_list),
+        "Edges": pd.DataFrame(edge_list),
+    }
     if edge_cls_list:
         tables["EdgesCondClass"] = pd.DataFrame(edge_cls_list)
     for i in node_dict:
@@ -80,4 +86,6 @@ def to_excel(structure, structure_name, dir="structure_storage"):
     node_list, node_dict = output_node_list(structure)
     edge_list, edge_cls_list = output_edge_list(structure)
 
-    _output_excel(node_list, node_dict, edge_list, edge_cls_list, structure_name, dir)
+    _output_excel(
+        node_list, node_dict, edge_list, edge_cls_list, structure_name, dir
+    )

lionagi/libs/ln_api.py

@@ -3,9 +3,9 @@ import contextlib
 import logging
 import re
 from abc import ABC
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from dataclasses import dataclass
-from typing import Any, Callable, NoReturn, Type
+from typing import Any, NoReturn, Type
 
 import aiohttp
 
@@ -70,7 +70,9 @@ class APIUtil:
                 False
         """
         if "error" in response_json:
-            logging.warning(f"API call failed with error: {response_json['error']}")
+            logging.warning(
+                f"API call failed with error: {response_json['error']}"
+            )
             return True
         return False
 
@@ -93,7 +95,9 @@ class APIUtil:
                 >>> api_rate_limit_error(response_json_without_rate_limit)
                 False
         """
-        return "Rate limit" in response_json.get("error", {}).get("message", "")
+        return "Rate limit" in response_json.get("error", {}).get(
+            "message", ""
+        )
 
     @staticmethod
     @func_call.lru_cache(maxsize=128)
@@ -389,15 +393,15 @@ class APIUtil:
                                     num_tokens += ImageUtil.calculate_image_token_usage_from_base64(
                                         a, item.get("detail", "low")
                                     )
-                                    num_tokens += (
-                                        20  # for every image we add 20 tokens buffer
-                                    )
+                                    num_tokens += 20  # for every image we add 20 tokens buffer
                             elif isinstance(item, str):
                                 num_tokens += len(encoding.encode(item))
                             else:
                                 num_tokens += len(encoding.encode(str(item)))
 
-                num_tokens += 2  # every reply is primed with <im_start>assistant
+                num_tokens += (
+                    2  # every reply is primed with <im_start>assistant
+                )
                 return num_tokens + completion_tokens
             else:
                 prompt = payload["format_prompt"]
@@ -405,7 +409,9 @@ class APIUtil:
                     prompt_tokens = len(encoding.encode(prompt))
                     return prompt_tokens + completion_tokens
                 elif isinstance(prompt, list):  # multiple prompts
-                    prompt_tokens = sum(len(encoding.encode(p)) for p in prompt)
+                    prompt_tokens = sum(
+                        len(encoding.encode(p)) for p in prompt
+                    )
                     return prompt_tokens + completion_tokens * len(prompt)
                 else:
                     raise TypeError(
@@ -427,7 +433,9 @@ class APIUtil:
             )
 
     @staticmethod
-    def create_payload(input_, config, required_, optional_, input_key, **kwargs):
+    def create_payload(
+        input_, config, required_, optional_, input_key, **kwargs
+    ):
         config = {**config, **kwargs}
         payload = {input_key: input_}
 
@@ -514,7 +522,9 @@ class BaseRateLimiter(ABC):
         except asyncio.CancelledError:
             logging.info("Rate limit replenisher task cancelled.")
         except Exception as e:
-            logging.error(f"An error occurred in the rate limit replenisher: {e}")
+            logging.error(
+                f"An error occurred in the rate limit replenisher: {e}"
+            )
 
     async def stop_replenishing(self) -> None:
         """Stops the replenishment task."""
@@ -658,7 +668,9 @@ class SimpleRateLimiter(BaseRateLimiter):
         token_encoding_name=None,
     ) -> None:
         """Initializes the SimpleRateLimiter with the specified parameters."""
-        super().__init__(max_requests, max_tokens, interval, token_encoding_name)
+        super().__init__(
+            max_requests, max_tokens, interval, token_encoding_name
+        )
 
 
 class EndPoint:
@@ -695,7 +707,7 @@ class EndPoint:
         max_tokens: int = 100000,
         interval: int = 60,
         endpoint_: str | None = None,
-        rate_limiter_class: Type[BaseRateLimiter] = SimpleRateLimiter,
+        rate_limiter_class: type[BaseRateLimiter] = SimpleRateLimiter,
         token_encoding_name=None,
         config: Mapping = None,
     ) -> None:
@@ -712,7 +724,10 @@ class EndPoint:
     async def init_rate_limiter(self) -> None:
         """Initializes the rate limiter for the endpoint."""
         self.rate_limiter = await self.rate_limiter_class.create(
-            self.max_requests, self.max_tokens, self.interval, self.token_encoding_name
+            self.max_requests,
+            self.max_tokens,
+            self.interval,
+            self.token_encoding_name,
         )
         self._has_initialized = True
 
@@ -786,7 +801,9 @@ class BaseService:
                             max_tokens=self.chat_config_rate_limit.get(
                                 "max_tokens", 100000
                             ),
-                            interval=self.chat_config_rate_limit.get("interval", 60),
+                            interval=self.chat_config_rate_limit.get(
+                                "interval", 60
+                            ),
                             endpoint_=ep,
                             token_encoding_name=self.token_encoding_name,
                             config=endpoint_config,
@@ -795,17 +812,20 @@ class BaseService:
                         self.endpoints[ep] = EndPoint(
                             max_requests=(
                                 endpoint_config.get("max_requests", 1000)
-                                if endpoint_config.get("max_requests", 1000) is not None
+                                if endpoint_config.get("max_requests", 1000)
+                                is not None
                                 else 1000
                             ),
                             max_tokens=(
                                 endpoint_config.get("max_tokens", 100000)
-                                if endpoint_config.get("max_tokens", 100000) is not None
+                                if endpoint_config.get("max_tokens", 100000)
+                                is not None
                                 else 100000
                             ),
                             interval=(
                                 endpoint_config.get("interval", 60)
-                                if endpoint_config.get("interval", 60) is not None
+                                if endpoint_config.get("interval", 60)
+                                is not None
                                 else 60
                             ),
                             endpoint_=ep,
@@ -832,7 +852,9 @@ class BaseService:
                 if not self.endpoints[ep]._has_initialized:
                     await self.endpoints[ep].init_rate_limiter()
 
-    async def call_api(self, payload, endpoint, method, required_tokens=None, **kwargs):
+    async def call_api(
+        self, payload, endpoint, method, required_tokens=None, **kwargs
+    ):
         """
         Calls the specified API endpoint with the given payload and method.
 

lionagi/libs/ln_context.py

@@ -26,12 +26,12 @@ async def async_suppress_print():
     """
     An asynchronous context manager that redirects stdout to /dev/null to suppress print output.
     """
-    original_stdout = sys.stdout  # Save the reference to the original standard output
+    original_stdout = (
+        sys.stdout
+    )  # Save the reference to the original standard output
     with open(os.devnull, "w") as devnull:
         sys.stdout = devnull
         try:
             yield
         finally:
-            sys.stdout = (
-                original_stdout  # Restore standard output to the original value
-            )
+            sys.stdout = original_stdout  # Restore standard output to the original value

lionagi/libs/ln_convert.py

@@ -1,7 +1,8 @@
 import json
 import re
+from collections.abc import Generator, Iterable
 from functools import singledispatch
-from typing import Any, Generator, Iterable, Type
+from typing import Any, Type
 
 import pandas as pd
 from pydantic import BaseModel
@@ -11,7 +12,9 @@ number_regex = re.compile(r"-?\d+\.?\d*")
 
 # to_list functions with datatype overloads
 @singledispatch
-def to_list(input_, /, *, flatten: bool = True, dropna: bool = True) -> list[Any]:
+def to_list(
+    input_, /, *, flatten: bool = True, dropna: bool = True
+) -> list[Any]:
     """
     Converts the input object to a list. This function is capable of handling various input types,
     utilizing single dispatch to specialize for different types such as list, tuple, and set.
@@ -43,9 +46,13 @@ def to_list(input_, /, *, flatten: bool = True, dropna: bool = True) -> list[Any
         ):
             return [input_]
         iterable_list = list(input_)
-        return _flatten_list(iterable_list, dropna) if flatten else iterable_list
+        return (
+            _flatten_list(iterable_list, dropna) if flatten else iterable_list
+        )
     except Exception as e:
-        raise ValueError(f"Could not convert {type(input_)} object to list: {e}") from e
+        raise ValueError(
+            f"Could not convert {type(input_)} object to list: {e}"
+        ) from e
 
 
 @to_list.register(list)
@@ -362,7 +369,9 @@ def _(
 ) -> pd.DataFrame:
     if not input_:
         return pd.DataFrame()
-    if not isinstance(input_[0], (pd.DataFrame, pd.Series, pd.core.generic.NDFrame)):
+    if not isinstance(
+        input_[0], (pd.DataFrame, pd.Series, pd.core.generic.NDFrame)
+    ):
         if drop_kwargs is None:
             drop_kwargs = {}
         try:
@@ -370,7 +379,9 @@ def _(
             dfs = dfs.dropna(**(drop_kwargs | {"how": how}))
             return dfs.reset_index(drop=True) if reset_index else dfs
         except Exception as e:
-            raise ValueError(f"Error converting input_ to DataFrame: {e}") from e
+            raise ValueError(
+                f"Error converting input_ to DataFrame: {e}"
+            ) from e
 
     dfs = ""
     if drop_kwargs is None:
@@ -401,7 +412,7 @@ def to_num(
     *,
     upper_bound: int | float | None = None,
     lower_bound: int | float | None = None,
-    num_type: Type[int | float] = float,
+    num_type: type[int | float] = float,
     precision: int | None = None,
 ) -> int | float:
     """
@@ -431,13 +442,17 @@ def to_readable_dict(input_: Any) -> str:
 
     try:
         dict_ = to_dict(input_)
-        return json.dumps(dict_, indent=4) if isinstance(input_, dict) else input_
+        return (
+            json.dumps(dict_, indent=4) if isinstance(input_, dict) else input_
+        )
     except Exception as e:
-        raise ValueError(f"Could not convert given input to readable dict: {e}") from e
+        raise ValueError(
+            f"Could not convert given input to readable dict: {e}"
+        ) from e
 
 
 def is_same_dtype(
-    input_: list | dict, dtype: Type | None = None, return_dtype=False
+    input_: list | dict, dtype: type | None = None, return_dtype=False
 ) -> bool:
     """
     Checks if all elements in a list or dictionary values are of the same data type.
@@ -496,7 +511,9 @@ def strip_lower(input_: Any) -> str:
     try:
         return str(input_).strip().lower()
     except Exception as e:
-        raise ValueError(f"Could not convert input_ to string: {input_}, Error: {e}")
+        raise ValueError(
+            f"Could not convert input_ to string: {input_}, Error: {e}"
+        )
 
 
 def is_structure_homogeneous(
@@ -552,7 +569,9 @@ def is_structure_homogeneous(
     return (is_, structure_type) if return_structure_type else is_
 
 
-def is_homogeneous(iterables: list[Any] | dict[Any, Any], type_check: type) -> bool:
+def is_homogeneous(
+    iterables: list[Any] | dict[Any, Any], type_check: type
+) -> bool:
     if isinstance(iterables, list):
         return all(isinstance(it, type_check) for it in iterables)
     return isinstance(iterables, type_check)
@@ -562,7 +581,7 @@ def _str_to_num(
     input_: str,
     upper_bound: float | None = None,
     lower_bound: float | None = None,
-    num_type: Type[int | float] = int,
+    num_type: type[int | float] = int,
     precision: int | None = None,
 ) -> int | float:
     number_str = _extract_first_number(input_)
@@ -590,7 +609,9 @@ def _extract_first_number(input_: str) -> str | None:
 
 
 def _convert_to_num(
-    number_str: str, num_type: Type[int | float] = int, precision: int | None = None
+    number_str: str,
+    num_type: type[int | float] = int,
+    precision: int | None = None,
 ) -> int | float:
     if num_type is int:
         return int(float(number_str))

lionagi/libs/ln_dataframe.py

@@ -31,7 +31,9 @@ def extend_dataframe(
     try:
         if len(df2.dropna(how="all")) > 0 and len(df1.dropna(how="all")) > 0:
             df = convert.to_df([df1, df2])
-            df.drop_duplicates(inplace=True, subset=[unique_col], keep=keep, **kwargs)
+            df.drop_duplicates(
+                inplace=True, subset=[unique_col], keep=keep, **kwargs
+            )
             df_ = convert.to_df(df)
             if len(df_) > 1:
                 return df_
@@ -112,7 +114,9 @@ def replace_keyword(
             keyword, replacement, case=False, regex=False
         )
     else:
-        df_.loc[:, column] = df_[column].str.replace(keyword, replacement, regex=False)
+        df_.loc[:, column] = df_[column].str.replace(
+            keyword, replacement, regex=False
+        )
 
     return df_ if inplace else True
 
@@ -160,7 +164,9 @@ def remove_last_n_rows(df: pd.DataFrame, steps: int) -> pd.DataFrame:
     return convert.to_df(df[:-steps])
 
 
-def update_row(df: pd.DataFrame, row: str | int, column: str | int, value: Any) -> bool:
+def update_row(
+    df: pd.DataFrame, row: str | int, column: str | int, value: Any
+) -> bool:
     """
     Updates a row's value for a specified column in a DataFrame.
 

lionagi/libs/ln_func_call.py

@@ -3,8 +3,9 @@ from __future__ import annotations
 import asyncio
 import functools
 import logging
+from collections.abc import Callable, Coroutine
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Callable, Coroutine
+from typing import Any
 
 from lionagi.libs.ln_async import AsyncUtil
 from lionagi.libs.ln_convert import to_list
@@ -65,9 +66,13 @@ def lcall(
     """
     lst = to_list(input_, dropna=dropna)
     if len(to_list(func)) != 1:
-        raise ValueError("There must be one and only one function for list calling.")
+        raise ValueError(
+            "There must be one and only one function for list calling."
+        )
 
-    return to_list([func(i, **kwargs) for i in lst], flatten=flatten, dropna=dropna)
+    return to_list(
+        [func(i, **kwargs) for i in lst], flatten=flatten, dropna=dropna
+    )
 
 
 async def alcall(
@@ -122,7 +127,9 @@ async def alcall(
     outs = await asyncio.gather(*tasks)
     outs_ = []
     for i in outs:
-        outs_.append(await i if isinstance(i, (Coroutine, asyncio.Future)) else i)
+        outs_.append(
+            await i if isinstance(i, (Coroutine, asyncio.Future)) else i
+        )
 
     return to_list(outs_, flatten=flatten, dropna=dropna)
 
@@ -260,7 +267,9 @@ async def tcall(
     async def async_call() -> tuple[Any, float]:
         start_time = SysUtil.get_now(datetime_=False)
         if timeout is not None:
-            result = await AsyncUtil.execute_timeout(func(*args, **kwargs), timeout)
+            result = await AsyncUtil.execute_timeout(
+                func(*args, **kwargs), timeout
+            )
             duration = SysUtil.get_now(datetime_=False) - start_time
             return (result, duration) if timing else result
         try:
@@ -282,12 +291,18 @@ async def tcall(
             handle_error(e)
 
     def handle_error(e: Exception):
-        _msg = f"{err_msg} Error: {e}" if err_msg else f"An error occurred: {e}"
+        _msg = (
+            f"{err_msg} Error: {e}" if err_msg else f"An error occurred: {e}"
+        )
         print(_msg)
         if not ignore_err:
             raise
 
-    return await async_call() if AsyncUtil.is_coroutine_func(func) else sync_call()
+    return (
+        await async_call()
+        if AsyncUtil.is_coroutine_func(func)
+        else sync_call()
+    )
 
 
 async def rcall(
@@ -348,7 +363,9 @@ async def rcall(
     start = SysUtil.get_now(datetime_=False)
     for attempt in range(retries + 1) if retries == 0 else range(retries):
         try:
-            err_msg = f"Attempt {attempt + 1}/{retries}: " if retries > 0 else None
+            err_msg = (
+                f"Attempt {attempt + 1}/{retries}: " if retries > 0 else None
+            )
             if timing:
                 return (
                     await _tcall(
@@ -362,7 +379,9 @@ async def rcall(
             last_exception = e
             if attempt < retries:
                 if verbose:
-                    print(f"Attempt {attempt + 1}/{retries} failed: {e}, retrying...")
+                    print(
+                        f"Attempt {attempt + 1}/{retries} failed: {e}, retrying..."
+                    )
                 await asyncio.sleep(delay)
                 delay *= backoff_factor
             else:
@@ -485,9 +504,13 @@ async def _tcall(
                 else default
             )
         else:
-            raise asyncio.TimeoutError(err_msg)  # Re-raise the timeout exception
+            raise asyncio.TimeoutError(
+                err_msg
+            )  # Re-raise the timeout exception
     except Exception as e:
-        err_msg = f"{err_msg} Error: {e}" if err_msg else f"An error occurred: {e}"
+        err_msg = (
+            f"{err_msg} Error: {e}" if err_msg else f"An error occurred: {e}"
+        )
         if ignore_err:
             return (
                 (default, SysUtil.get_now(datetime_=False) - start_time)
@@ -645,7 +668,9 @@ class CallDecorator:
         def decorator(func: Callable[..., Any]) -> Callable:
             @functools.wraps(func)
             async def wrapper(*args, **kwargs) -> Any:
-                return await rcall(func, *args, default=default_value, **kwargs)
+                return await rcall(
+                    func, *args, default=default_value, **kwargs
+                )
 
             return wrapper
 
@@ -863,8 +888,12 @@ class CallDecorator:
                         k: preprocess(v, *preprocess_args, **preprocess_kwargs)
                         for k, v in kwargs.items()
                     }
-                    result = await func(*preprocessed_args, **preprocessed_kwargs)
-                    return postprocess(result, *postprocess_args, **postprocess_kwargs)
+                    result = await func(
+                        *preprocessed_args, **preprocessed_kwargs
+                    )
+                    return postprocess(
+                        result, *postprocess_args, **postprocess_kwargs
+                    )
 
                 return async_wrapper
             else:
@@ -880,7 +909,9 @@ class CallDecorator:
                         for k, v in kwargs.items()
                     }
                     result = func(*preprocessed_args, **preprocessed_kwargs)
-                    return postprocess(result, *postprocess_args, **postprocess_kwargs)
+                    return postprocess(
+                        result, *postprocess_args, **postprocess_kwargs
+                    )
 
                 return sync_wrapper
 
@@ -1165,7 +1196,9 @@ class Throttle:
 
         return wrapper
 
-    async def __call_async__(self, func: Callable[..., Any]) -> Callable[..., Any]:
+    async def __call_async__(
+        self, func: Callable[..., Any]
+    ) -> Callable[..., Any]:
         """
         Decorates an asynchronous function with the throttling mechanism.
 
@@ -1187,7 +1220,9 @@ class Throttle:
         return wrapper
 
 
-def _custom_error_handler(error: Exception, error_map: dict[type, Callable]) -> None:
+def _custom_error_handler(
+    error: Exception, error_map: dict[type, Callable]
+) -> None:
     # noinspection PyUnresolvedReferences
     """
     handle errors based on a given error mapping.
@@ -1262,7 +1297,7 @@ async def call_handler(
         raise
 
 
-@functools.lru_cache(maxsize=None)
+@functools.cache
 def is_coroutine_func(func: Callable) -> bool:
     """
     checks if the specified function is an asyncio coroutine function.

lionagi/libs/ln_image.py

@@ -10,7 +10,7 @@ class ImageUtil:
 
     @staticmethod
     def preprocess_image(
-        image: np.ndarray, color_conversion_code: Optional[int] = None
+        image: np.ndarray, color_conversion_code: int | None = None
     ) -> np.ndarray:
         SysUtil.check_import("cv2", pip_name="opencv-python")
         import cv2
@@ -19,19 +19,23 @@ class ImageUtil:
         return cv2.cvtColor(image, color_conversion_code)
 
     @staticmethod
-    def encode_image_to_base64(image: np.ndarray, file_extension: str = ".jpg") -> str:
+    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.")
+            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
+        image_path: str, color_flag: int | None = None
     ) -> np.ndarray:
         SysUtil.check_import("cv2", pip_name="opencv-python")
         import cv2
@@ -45,7 +49,7 @@ class ImageUtil:
     @staticmethod
     def read_image_to_base64(
         image_path: str,
-        color_flag: Optional[int] = None,
+        color_flag: int | None = None,
     ) -> str:
         image_path = str(image_path)
         image = ImageUtil.read_image_to_array(image_path, color_flag)

lionagi/libs/ln_knowledge_graph.py

@@ -33,7 +33,9 @@ class KnowledgeBase:
         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,
+        self.relations = (
+            []
+        )  # [ head: entity_title, type: ..., tail: entity_title,
         # meta: { article_url: { spans: [...] } } ]
         self.sources = {}  # { article_url: {...} }
 
@@ -48,7 +50,9 @@ class KnowledgeBase:
             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"]
+                r,
+                source_data["article_title"],
+                source_data["article_publish_date"],
             )
 
     def are_relations_equal(self, r1, r2):
@@ -137,7 +141,9 @@ class KnowledgeBase:
         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"}
+        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):
         """
@@ -210,7 +216,9 @@ class KnowledgeBase:
         relation, subject, relation, object_ = "", "", "", ""
         text = text.strip()
         current = "x"
-        text_replaced = text.replace("<s>", "").replace("<pad>", "").replace("</s>", "")
+        text_replaced = (
+            text.replace("<s>", "").replace("<pad>", "").replace("</s>", "")
+        )
         for token in text_replaced.split():
             if token == "<triplet>":
                 current = "t"
@@ -311,7 +319,9 @@ class KGTripletExtractor:
 
         if not any([model, tokenizer]):
             tokenizer = AutoTokenizer.from_pretrained("Babelscape/rebel-large")
-            model = AutoModelForSeq2SeqLM.from_pretrained("Babelscape/rebel-large")
+            model = AutoModelForSeq2SeqLM.from_pretrained(
+                "Babelscape/rebel-large"
+            )
             model.to(device)
 
         inputs = tokenizer([text], return_tensors="pt")
@@ -373,10 +383,14 @@ class KGTripletExtractor:
         i = 0
         for sentence_pred in decoded_preds:
             current_span_index = i // num_return_sequences
-            relations = KnowledgeBase.extract_relations_from_model_output(sentence_pred)
+            relations = KnowledgeBase.extract_relations_from_model_output(
+                sentence_pred
+            )
             for relation in relations:
                 relation["meta"] = {
-                    "article_url": {"spans": [spans_boundaries[current_span_index]]}
+                    "article_url": {
+                        "spans": [spans_boundaries[current_span_index]]
+                    }
                 }
                 kb.add_relation(relation, article_title, article_publish_date)
             i += 1