From e62c665a82f3f41ffb7986483ff37b9dc30f2d27 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Thu, 4 Jan 2024 11:50:30 -0500
Subject: [PATCH 1/8] [FEAT][Modelscope] [CogAgent] [REMOVAL][Vllm] [v]

---
 pyproject.toml              |   2 +-
 requirements.txt            |   1 +
 swarms/agents/tool_agent.py |   5 +-
 swarms/models/__init__.py   |   4 +-
 swarms/models/cog_agent.py  | 128 ++++++++++++++++++++++++++++++++++++
 5 files changed, 135 insertions(+), 5 deletions(-)
 create mode 100644 swarms/models/cog_agent.py

diff --git a/pyproject.toml b/pyproject.toml
index 0871ddaa..f0551c4c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -68,9 +68,9 @@ rich = "13.5.2"
 sqlalchemy = "*"
 pgvector = "*"
 qdrant-client = "*"
-vllm = "*"
 sentence-transformers = "*"
 peft = "*"
+modelscope = "1.10.0"
 
 
 [tool.poetry.group.lint.dependencies]
diff --git a/requirements.txt b/requirements.txt
index 9944b616..238b32fb 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -75,3 +75,4 @@ qdrant-client
 vllm
 sentence-transformers
 peft
+modelscope==1.10.0
\ No newline at end of file
diff --git a/swarms/agents/tool_agent.py b/swarms/agents/tool_agent.py
index 594a1863..bc34a476 100644
--- a/swarms/agents/tool_agent.py
+++ b/swarms/agents/tool_agent.py
@@ -61,6 +61,7 @@ class ToolAgent(AbstractLLM):
 
         print(generated_data)
     """
+
     def __init__(
         self,
         name: str,
@@ -108,7 +109,7 @@ class ToolAgent(AbstractLLM):
         except Exception as error:
             print(f"[Error] [ToolAgent] {error}")
             raise error
-    
+
     def __call__(self, task: str, *args, **kwargs):
         """Call self as a function.
 
@@ -118,4 +119,4 @@ class ToolAgent(AbstractLLM):
         Returns:
             _type_: _description_
         """
-        return self.run(task, *args, **kwargs)
\ No newline at end of file
+        return self.run(task, *args, **kwargs)
diff --git a/swarms/models/__init__.py b/swarms/models/__init__.py
index 58701f64..640a7297 100644
--- a/swarms/models/__init__.py
+++ b/swarms/models/__init__.py
@@ -9,7 +9,7 @@ from swarms.models.openai_models import (
     OpenAIChat,
 )  # noqa: E402
 
-from swarms.models.vllm import vLLM  # noqa: E402
+# from swarms.models.vllm import vLLM  # noqa: E402
 from swarms.models.zephyr import Zephyr  # noqa: E402
 from swarms.models.biogpt import BioGPT  # noqa: E402
 from swarms.models.huggingface import HuggingfaceLLM  # noqa: E402
@@ -72,7 +72,7 @@ __all__ = [
     # "Dalle3",
     # "DistilWhisperModel",
     "GPT4VisionAPI",
-    "vLLM",
+    # "vLLM",
     "OpenAITTS",
     "Gemini",
     "Gigabind",
diff --git a/swarms/models/cog_agent.py b/swarms/models/cog_agent.py
new file mode 100644
index 00000000..7a3ff684
--- /dev/null
+++ b/swarms/models/cog_agent.py
@@ -0,0 +1,128 @@
+import torch
+from PIL import Image
+from modelscope import AutoModelForCausalLM, AutoTokenizer
+from swarms.models.base_multimodal_model import BaseMultiModalModel
+
+device_check = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class CogAgent(BaseMultiModalModel):
+    """CogAgent
+    
+    Multi-modal conversational agent that can be used to chat with
+    images and text. It is based on the CogAgent model from the
+    ModelScope library.
+    
+    Attributes:
+        model_name (str): The name of the model to be used
+        tokenizer_name (str): The name of the tokenizer to be used
+        dtype (torch.bfloat16): The data type to be used
+        low_cpu_mem_usage (bool): Whether to use low CPU memory
+        load_in_4bit (bool): Whether to load in 4-bit
+        trust_remote_code (bool): Whether to trust remote code
+        device (str): The device to be used
+    
+    Examples:
+        >>> from swarms.models.cog_agent import CogAgent
+        >>> cog_agent = CogAgent()
+        >>> cog_agent.run("How are you?", "images/1.jpg")
+        <s> I'm fine. How are you? </s>
+    """
+    def __init__(
+        self,
+        model_name: str = "ZhipuAI/cogagent-chat",
+        tokenizer_name: str = "I-ModelScope/vicuna-7b-v1.5",
+        dtype=torch.bfloat16,
+        low_cpu_mem_usage: bool = True,
+        load_in_4bit: bool = True,
+        trust_remote_code: bool = True,
+        device=device_check,
+        *args,
+        **kwargs,
+    ):
+        super().__init__()
+        self.model_name = model_name
+        self.tokenizer_name = tokenizer_name
+        self.dtype = dtype
+        self.low_cpu_mem_usage = low_cpu_mem_usage
+        self.load_in_4bit = load_in_4bit
+        self.trust_remote_code = trust_remote_code
+        self.device = device
+
+        self.model = (
+            AutoModelForCausalLM.from_pretrained(
+                self.model_name,
+                torch_dtype=self.dtype,
+                low_cpu_mem_usage=self.low_cpu_mem_usage,
+                load_in_4bit=self.load_in_4bit,
+                trust_remote_code=self.trust_remote_code,
+                *args,
+                **kwargs,
+            )
+            .to(self.device)
+            .eval()
+        )
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.tokenizer_name
+        )
+
+    def run(self, task: str, img: str, *args, **kwargs):
+        """Run the model
+
+        Args:
+            task (str): The task to be performed
+            img (str): The image path
+            
+        """ 
+        image = Image.open(img).convert("RGB")
+
+        input_by_model = self.model.build_conversation_input_ids(
+            self.tokenizer,
+            query=task,
+            history=[],
+            images=[image],
+        )
+
+        inputs = {
+            "input_ids": (
+                input_by_model["input_ids"]
+                .unsqueeze(0)
+                .to(self.device)
+            ),
+            "token_type_ids": (
+                input_by_model["token_type_ids"]
+                .unsqueeze(0)
+                .to(self.device)
+            ),
+            "attention_mask": (
+                input_by_model["attention_mask"]
+                .unsqueeze(0)
+                .to(self.device)
+            ),
+            "images": [
+                [
+                    input_by_model["images"][0]
+                    .to(self.device)
+                    .to(self.dtype)
+                ]
+            ],
+        }
+        if (
+            "cross_images" in input_by_model
+            and input_by_model["cross_images"]
+        ):
+            inputs["cross_images"] = [
+                [
+                    input_by_model["cross_images"][0]
+                    .to(self.device)
+                    .to(self.dtype)
+                ]
+            ]
+
+        with torch.no_grad():
+            outputs = self.model(**inputs, **kwargs)
+            outputs = outputs[:, inputs["input_ids"].shape[1] :]
+            response = self.decode(outputs[0])
+            response = response.split("</s>")[0]
+            print(response)

From 35e62368c04d6087a9ec6a39fd418813733da568 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Thu, 4 Jan 2024 11:51:55 -0500
Subject: [PATCH 2/8] [REQUIREMENTS][Vllm][removal]

---
 pyproject.toml   | 2 +-
 requirements.txt | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index f0551c4c..847470eb 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms"
-version = "3.2.5"
+version = "3.2.6"
 description = "Swarms - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <kye@apac.ai>"]
diff --git a/requirements.txt b/requirements.txt
index 238b32fb..9250997a 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -72,7 +72,6 @@ pre-commit==3.2.2
 sqlalchemy
 pgvector
 qdrant-client
-vllm
 sentence-transformers
 peft
 modelscope==1.10.0
\ No newline at end of file

From 41b858a91d6262b49810aa39cd8e9c094a568775 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Thu, 4 Jan 2024 11:53:24 -0500
Subject: [PATCH 3/8] [CLEANUP][swarms.models.__init__]

---
 pyproject.toml            | 2 +-
 swarms/models/__init__.py | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 847470eb..a40d9530 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms"
-version = "3.2.6"
+version = "3.2.7"
 description = "Swarms - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <kye@apac.ai>"]
diff --git a/swarms/models/__init__.py b/swarms/models/__init__.py
index 640a7297..12436adf 100644
--- a/swarms/models/__init__.py
+++ b/swarms/models/__init__.py
@@ -40,6 +40,7 @@ from swarms.models.zeroscope import ZeroscopeTTV  # noqa: E402
 # from swarms.models.distilled_whisperx import DistilWhisperModel # noqa: E402
 # from swarms.models.whisperx_model import WhisperX  # noqa: E402
 # from swarms.models.kosmos_two import Kosmos  # noqa: E402
+# from swarms.models.cog_agent import CogAgent  # noqa: E402
 
 from swarms.models.types import (
     TextModality,
@@ -83,4 +84,5 @@ __all__ = [
     "AudioModality",
     "VideoModality",
     "MultimodalData",
+    # "CogAgent"
 ]

From e8ca14f071d765e2fdcc5d53d72ac592bbe1d810 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Fri, 5 Jan 2024 12:29:24 -0500
Subject: [PATCH 4/8] [REFACTOR][HuggingfaceLLM]

---
 pyproject.toml               |   2 +-
 swarms/models/cog_agent.py   |  11 +--
 swarms/models/huggingface.py | 137 +++++++----------------------------
 3 files changed, 33 insertions(+), 117 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index a40d9530..0fa7b2db 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms"
-version = "3.2.7"
+version = "3.2.8"
 description = "Swarms - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <kye@apac.ai>"]
diff --git a/swarms/models/cog_agent.py b/swarms/models/cog_agent.py
index 7a3ff684..2d0d09e9 100644
--- a/swarms/models/cog_agent.py
+++ b/swarms/models/cog_agent.py
@@ -8,11 +8,11 @@ device_check = "cuda" if torch.cuda.is_available() else "cpu"
 
 class CogAgent(BaseMultiModalModel):
     """CogAgent
-    
+
     Multi-modal conversational agent that can be used to chat with
     images and text. It is based on the CogAgent model from the
     ModelScope library.
-    
+
     Attributes:
         model_name (str): The name of the model to be used
         tokenizer_name (str): The name of the tokenizer to be used
@@ -21,13 +21,14 @@ class CogAgent(BaseMultiModalModel):
         load_in_4bit (bool): Whether to load in 4-bit
         trust_remote_code (bool): Whether to trust remote code
         device (str): The device to be used
-    
+
     Examples:
         >>> from swarms.models.cog_agent import CogAgent
         >>> cog_agent = CogAgent()
         >>> cog_agent.run("How are you?", "images/1.jpg")
         <s> I'm fine. How are you? </s>
     """
+
     def __init__(
         self,
         model_name: str = "ZhipuAI/cogagent-chat",
@@ -73,8 +74,8 @@ class CogAgent(BaseMultiModalModel):
         Args:
             task (str): The task to be performed
             img (str): The image path
-            
-        """ 
+
+        """
         image = Image.open(img).convert("RGB")
 
         input_by_model = self.model.build_conversation_input_ids(
diff --git a/swarms/models/huggingface.py b/swarms/models/huggingface.py
index bbb39223..d9447f3c 100644
--- a/swarms/models/huggingface.py
+++ b/swarms/models/huggingface.py
@@ -3,18 +3,18 @@ import concurrent.futures
 import logging
 from typing import List, Tuple
 
-
 import torch
 from termcolor import colored
-from torch.nn.parallel import DistributedDataParallel as DDP
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     BitsAndBytesConfig,
 )
 
+from swarms.models.base_llm import AbstractLLM
+
 
-class HuggingfaceLLM:
+class HuggingfaceLLM(AbstractLLM):
     """
     A class for running inference on a given model.
 
@@ -123,7 +123,6 @@ class HuggingfaceLLM:
         quantize: bool = False,
         quantization_config: dict = None,
         verbose=False,
-        # logger=None,
         distributed=False,
         decoding=False,
         max_workers: int = 5,
@@ -135,6 +134,7 @@ class HuggingfaceLLM:
         *args,
         **kwargs,
     ):
+        super().__init__(*args, **kwargs)
         self.logger = logging.getLogger(__name__)
         self.device = (
             device
@@ -174,16 +174,21 @@ class HuggingfaceLLM:
 
         try:
             self.tokenizer = AutoTokenizer.from_pretrained(
-                self.model_id, *args, **kwargs
-            )
-            self.model = AutoModelForCausalLM.from_pretrained(
-                self.model_id,
-                quantization_config=bnb_config,
-                *args,
-                **kwargs,
+                self.model_id
             )
 
-            self.model  # .to(self.device)
+            if quantize:
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    self.model_id,
+                    quantization_config=bnb_config,
+                    *args,
+                    **kwargs,
+                ).to(self.device)
+            else:
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    self.model_id, *args, **kwargs
+                ).to(self.device)
+
         except Exception as e:
             # self.logger.error(f"Failed to load the model or the tokenizer: {e}")
             # raise
@@ -205,33 +210,6 @@ class HuggingfaceLLM:
         """Ashcnronous generate text for a given prompt"""
         return await asyncio.to_thread(self.run, task)
 
-    def load_model(self):
-        """Load the model"""
-        if not self.model or not self.tokenizer:
-            try:
-                self.tokenizer = AutoTokenizer.from_pretrained(
-                    self.model_id
-                )
-
-                bnb_config = (
-                    BitsAndBytesConfig(**self.quantization_config)
-                    if self.quantization_config
-                    else None
-                )
-
-                self.model = AutoModelForCausalLM.from_pretrained(
-                    self.model_id, quantization_config=bnb_config
-                ).to(self.device)
-
-                if self.distributed:
-                    self.model = DDP(self.model)
-            except Exception as error:
-                self.logger.error(
-                    "Failed to load the model or the tokenizer:"
-                    f" {error}"
-                )
-                raise
-
     def concurrent_run(self, tasks: List[str], max_workers: int = 5):
         """Concurrently generate text for a list of prompts."""
         with concurrent.futures.ThreadPoolExecutor(
@@ -252,7 +230,7 @@ class HuggingfaceLLM:
             results = [future.result() for future in futures]
         return results
 
-    def run(self, task: str):
+    def run(self, task: str, *args, **kwargs):
         """
         Generate a response based on the prompt text.
 
@@ -263,20 +241,12 @@ class HuggingfaceLLM:
         Returns:
         - Generated text (str).
         """
-        self.load_model()
-
-        max_length = self.max_length
-
-        self.print_dashboard(task)
-
         try:
             inputs = self.tokenizer.encode(task, return_tensors="pt")
 
-            # self.log.start()
-
             if self.decoding:
                 with torch.no_grad():
-                    for _ in range(max_length):
+                    for _ in range(self.max_length):
                         output_sequence = []
 
                         outputs = self.model.generate(
@@ -300,7 +270,11 @@ class HuggingfaceLLM:
             else:
                 with torch.no_grad():
                     outputs = self.model.generate(
-                        inputs, max_length=max_length, do_sample=True
+                        inputs,
+                        max_length=self.max_length,
+                        do_sample=True,
+                        *args,
+                        **kwargs,
                     )
 
             del inputs
@@ -320,67 +294,8 @@ class HuggingfaceLLM:
             )
             raise
 
-    def __call__(self, task: str):
-        """
-        Generate a response based on the prompt text.
-
-        Args:
-        - task (str): Text to prompt the model.
-        - max_length (int): Maximum length of the response.
-
-        Returns:
-        - Generated text (str).
-        """
-        self.load_model()
-
-        max_length = self.max_length
-
-        self.print_dashboard(task)
-
-        try:
-            inputs = self.tokenizer.encode(
-                task, return_tensors="pt"
-            ).to(self.device)
-
-            # self.log.start()
-
-            if self.decoding:
-                with torch.no_grad():
-                    for _ in range(max_length):
-                        output_sequence = []
-
-                        outputs = self.model.generate(
-                            inputs,
-                            max_length=len(inputs) + 1,
-                            do_sample=True,
-                        )
-                        output_tokens = outputs[0][-1]
-                        output_sequence.append(output_tokens.item())
-
-                        # print token in real-time
-                        print(
-                            self.tokenizer.decode(
-                                [output_tokens],
-                                skip_special_tokens=True,
-                            ),
-                            end="",
-                            flush=True,
-                        )
-                        inputs = outputs
-            else:
-                with torch.no_grad():
-                    outputs = self.model.generate(
-                        inputs, max_length=max_length, do_sample=True
-                    )
-
-            del inputs
-
-            return self.tokenizer.decode(
-                outputs[0], skip_special_tokens=True
-            )
-        except Exception as e:
-            self.logger.error(f"Failed to generate the text: {e}")
-            raise
+    def __call__(self, task: str, *args, **kwargs):
+        return self.run(task, *args, **kwargs)
 
     async def __call_async__(self, task: str, *args, **kwargs) -> str:
         """Call the model asynchronously""" ""

From 1f1135bb7dd62881bdcf6031fbe3230db12d3c5e Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Fri, 5 Jan 2024 13:21:49 -0500
Subject: [PATCH 5/8] [removal][vllm]

---
 pyproject.toml        | 8 ++++----
 requirements.txt      | 3 +--
 swarms/models/vllm.py | 4 +---
 3 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 0fa7b2db..c09edb2b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms"
-version = "3.2.8"
+version = "3.3.4"
 description = "Swarms - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <kye@apac.ai>"]
@@ -24,7 +24,7 @@ classifiers = [
 [tool.poetry.dependencies]
 python = "^3.6.1"
 torch = "2.1.1"
-transformers = "4.35.0"
+transformers = "4.36.2"
 openai = "0.28.0"
 langchain = "0.0.333"
 asyncio = "3.4.3"
@@ -37,7 +37,7 @@ opencv-python-headless = "4.8.1.78"
 faiss-cpu = "1.7.4"
 backoff = "2.2.1"
 marshmallow = "3.19.0"
-datasets = "2.10.1"
+datasets = "*"
 optimum = "1.15.0"
 diffusers = "*"
 PyPDF2 = "3.0.1"
@@ -53,7 +53,7 @@ ggl = "1.1.0"
 ratelimit = "2.2.1"
 beautifulsoup4 = "4.11.2"
 cohere = "4.24"
-huggingface-hub = "0.16.4"
+huggingface-hub = "*"
 pydantic = "1.10.12"
 tenacity = "8.2.2"
 Pillow = "9.4.0"
diff --git a/requirements.txt b/requirements.txt
index 9250997a..e8e2380e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
 torch==2.1.1
-transformers>2.10==4.35.0
+transformers>2.10==4.36.2
 pandas==1.5.3
 langchain==0.0.333
 nest_asyncio==1.5.6
@@ -72,6 +72,5 @@ pre-commit==3.2.2
 sqlalchemy
 pgvector
 qdrant-client
-sentence-transformers
 peft
 modelscope==1.10.0
\ No newline at end of file
diff --git a/swarms/models/vllm.py b/swarms/models/vllm.py
index 58745a75..0caeb3c8 100644
--- a/swarms/models/vllm.py
+++ b/swarms/models/vllm.py
@@ -1,13 +1,11 @@
 import torch
 from swarms.models.base_llm import AbstractLLM
-import subprocess
 
 if torch.cuda.is_available() or torch.cuda.device_count() > 0:
     # Download vllm with pip
     try:
-        subprocess.run(["pip", "install", "vllm"])
         from vllm import LLM, SamplingParams
-    except Exception as error:
+    except ImportError as error:
         print(f"[ERROR] [vLLM] {error}")
         raise error
 else:

From 5dd4758d1db7984be36de837fa9935747a0e0649 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Fri, 5 Jan 2024 14:50:55 -0500
Subject: [PATCH 6/8] [FEAT][BlockList]

---
 README.md                               | 109 ++++++++++++++++++++++++
 block.py                                |  97 +++++++++++++++++++++
 pyproject.toml                          |   5 +-
 requirements.txt                        |   1 +
 swarms/__init__.py                      |   4 +-
 swarms/models/huggingface.py            |  49 ++++-------
 swarms/structs/block_wrapper.py         |   4 +-
 swarms/structs/blockslist.py            |  33 ++++++-
 swarms/structs/model_parallizer.py      |  12 ---
 swarms/telemetry/auto_upgrade_swarms.py |   4 +
 swarms/telemetry/bootup.py              |   8 ++
 swarms/telemetry/check_update.py        |   4 -
 12 files changed, 277 insertions(+), 53 deletions(-)
 create mode 100644 block.py
 create mode 100644 swarms/telemetry/bootup.py

diff --git a/README.md b/README.md
index 2f9883a6..997bc797 100644
--- a/README.md
+++ b/README.md
@@ -456,6 +456,115 @@ print(f"Task result: {task.result}")
 ---
 
 
+
+### `BlockList`
+- Modularity and Flexibility: BlocksList allows users to create custom swarms by adding or removing different classes or functions as blocks. This means users can easily tailor the functionality of their swarm to suit their specific needs.
+
+- Ease of Management: With methods to add, remove, update, and retrieve blocks, BlocksList provides a straightforward way to manage the components of a swarm. This makes it easier to maintain and update the swarm over time.
+
+- Enhanced Searchability: BlocksList offers methods to get blocks by various attributes such as name, type, ID, and parent-related properties. This makes it easier for users to find and work with specific blocks in a large and complex swarm.
+
+```python
+import os
+
+from dotenv import load_dotenv
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Import the models, structs, and telemetry modules
+from swarms import (
+    Gemini,
+    GPT4VisionAPI,
+    Mixtral,
+    OpenAI,
+    ToolAgent,
+    BlocksList,
+)
+
+# Load the environment variables
+load_dotenv()
+
+# Get the environment variables
+openai_api_key = os.getenv("OPENAI_API_KEY")
+gemini_api_key = os.getenv("GEMINI_API_KEY")
+
+# Tool Agent
+model = AutoModelForCausalLM.from_pretrained(
+    "databricks/dolly-v2-12b"
+)
+tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")
+json_schema = {
+    "type": "object",
+    "properties": {
+        "name": {"type": "string"},
+        "age": {"type": "number"},
+        "is_student": {"type": "boolean"},
+        "courses": {"type": "array", "items": {"type": "string"}},
+    },
+}
+toolagent = ToolAgent(
+    model=model, tokenizer=tokenizer, json_schema=json_schema
+)
+
+# Blocks List which enables you to build custom swarms by adding classes or functions
+swarm = BlocksList(
+    "SocialMediaSwarm",
+    "A swarm of social media agents",
+    [
+        OpenAI(openai_api_key=openai_api_key),
+        Mixtral(),
+        GPT4VisionAPI(openai_api_key=openai_api_key),
+        Gemini(gemini_api_key=gemini_api_key),
+    ],
+)
+
+
+# Add the new block to the swarm
+swarm.add(toolagent)
+
+# Remove a block from the swarm
+swarm.remove(toolagent)
+
+# Update a block in the swarm
+swarm.update(toolagent)
+
+# Get a block at a specific index
+block_at_index = swarm.get(0)
+
+# Get all blocks in the swarm
+all_blocks = swarm.get_all()
+
+# Get blocks by name
+openai_blocks = swarm.get_by_name("OpenAI")
+
+# Get blocks by type
+gpt4_blocks = swarm.get_by_type("GPT4VisionAPI")
+
+# Get blocks by ID
+block_by_id = swarm.get_by_id(toolagent.id)
+
+# Get blocks by parent
+blocks_by_parent = swarm.get_by_parent(swarm)
+
+# Get blocks by parent ID
+blocks_by_parent_id = swarm.get_by_parent_id(swarm.id)
+
+# Get blocks by parent name
+blocks_by_parent_name = swarm.get_by_parent_name(swarm.name)
+
+# Get blocks by parent type
+blocks_by_parent_type = swarm.get_by_parent_type(type(swarm).__name__)
+
+# Get blocks by parent description
+blocks_by_parent_description = swarm.get_by_parent_description(
+    swarm.description
+)
+
+# Run the block in the swarm
+inference = swarm.run_block(toolagent, "Hello World")
+print(inference)
+```
+
+
 ## Real-World Deployment
 
 ### Multi-Agent Swarm for Logistics
diff --git a/block.py b/block.py
new file mode 100644
index 00000000..c5a76910
--- /dev/null
+++ b/block.py
@@ -0,0 +1,97 @@
+import os
+
+from dotenv import load_dotenv
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Import the models, structs, and telemetry modules
+from swarms import (
+    Gemini,
+    GPT4VisionAPI,
+    Mixtral,
+    OpenAI,
+    ToolAgent,
+    BlocksList,
+)
+
+# Load the environment variables
+load_dotenv()
+
+# Get the environment variables
+openai_api_key = os.getenv("OPENAI_API_KEY")
+gemini_api_key = os.getenv("GEMINI_API_KEY")
+
+# Tool Agent
+model = AutoModelForCausalLM.from_pretrained(
+    "databricks/dolly-v2-12b"
+)
+tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")
+json_schema = {
+    "type": "object",
+    "properties": {
+        "name": {"type": "string"},
+        "age": {"type": "number"},
+        "is_student": {"type": "boolean"},
+        "courses": {"type": "array", "items": {"type": "string"}},
+    },
+}
+toolagent = ToolAgent(
+    model=model, tokenizer=tokenizer, json_schema=json_schema
+)
+
+# Blocks List which enables you to build custom swarms by adding classes or functions
+swarm = BlocksList(
+    "SocialMediaSwarm",
+    "A swarm of social media agents",
+    [
+        OpenAI(openai_api_key=openai_api_key),
+        Mixtral(),
+        GPT4VisionAPI(openai_api_key=openai_api_key),
+        Gemini(gemini_api_key=gemini_api_key),
+    ],
+)
+
+
+# Add the new block to the swarm
+swarm.add(toolagent)
+
+# Remove a block from the swarm
+swarm.remove(toolagent)
+
+# Update a block in the swarm
+swarm.update(toolagent)
+
+# Get a block at a specific index
+block_at_index = swarm.get(0)
+
+# Get all blocks in the swarm
+all_blocks = swarm.get_all()
+
+# Get blocks by name
+openai_blocks = swarm.get_by_name("OpenAI")
+
+# Get blocks by type
+gpt4_blocks = swarm.get_by_type("GPT4VisionAPI")
+
+# Get blocks by ID
+block_by_id = swarm.get_by_id(toolagent.id)
+
+# Get blocks by parent
+blocks_by_parent = swarm.get_by_parent(swarm)
+
+# Get blocks by parent ID
+blocks_by_parent_id = swarm.get_by_parent_id(swarm.id)
+
+# Get blocks by parent name
+blocks_by_parent_name = swarm.get_by_parent_name(swarm.name)
+
+# Get blocks by parent type
+blocks_by_parent_type = swarm.get_by_parent_type(type(swarm).__name__)
+
+# Get blocks by parent description
+blocks_by_parent_description = swarm.get_by_parent_description(
+    swarm.description
+)
+
+# Run the block in the swarm
+inference = swarm.run_block(toolagent, "Hello World")
+print(inference)
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index c09edb2b..e0602d43 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms"
-version = "3.3.4"
+version = "3.4.1"
 description = "Swarms - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <kye@apac.ai>"]
@@ -41,7 +41,7 @@ datasets = "*"
 optimum = "1.15.0"
 diffusers = "*"
 PyPDF2 = "3.0.1"
-accelerate = "0.22.0"
+accelerate = "*"
 sentencepiece = "0.1.98"
 wget = "3.2"
 tensorflow = "2.14.0"
@@ -66,6 +66,7 @@ soundfile = "0.12.1"
 torchvision = "0.16.1"
 rich = "13.5.2"
 sqlalchemy = "*"
+bitsandbytes = "*"
 pgvector = "*"
 qdrant-client = "*"
 sentence-transformers = "*"
diff --git a/requirements.txt b/requirements.txt
index e8e2380e..eb39f71a 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -17,6 +17,7 @@ openai==0.28.0
 attrs==22.2.0
 datasets==2.10.1
 pydantic==1.10.12
+bitsandbytes
 soundfile==0.12.1
 arize-phoenix
 weaviate-client==3.25.3
diff --git a/swarms/__init__.py b/swarms/__init__.py
index 1c7514c4..4e6785cb 100644
--- a/swarms/__init__.py
+++ b/swarms/__init__.py
@@ -1,6 +1,6 @@
-from swarms.utils.disable_logging import disable_logging
+from swarms.telemetry.bootup import bootup  # noqa: E402, F403
 
-disable_logging()
+bootup()
 
 from swarms.agents import *  # noqa: E402, F403
 from swarms.structs import *  # noqa: E402, F403
diff --git a/swarms/models/huggingface.py b/swarms/models/huggingface.py
index d9447f3c..8e74cff5 100644
--- a/swarms/models/huggingface.py
+++ b/swarms/models/huggingface.py
@@ -131,6 +131,7 @@ class HuggingfaceLLM(AbstractLLM):
         temperature: float = 0.7,
         top_k: int = 40,
         top_p: float = 0.8,
+        dtype = torch.bfloat16,
         *args,
         **kwargs,
     ):
@@ -146,7 +147,6 @@ class HuggingfaceLLM(AbstractLLM):
         self.verbose = verbose
         self.distributed = distributed
         self.decoding = decoding
-        self.model, self.tokenizer = None, None
         self.quantize = quantize
         self.quantization_config = quantization_config
         self.max_workers = max_workers
@@ -155,6 +155,7 @@ class HuggingfaceLLM(AbstractLLM):
         self.temperature = temperature
         self.top_k = top_k
         self.top_p = top_p
+        self.dtype = dtype
 
         if self.distributed:
             assert (
@@ -168,39 +169,26 @@ class HuggingfaceLLM(AbstractLLM):
                     "load_in_4bit": True,
                     "bnb_4bit_use_double_quant": True,
                     "bnb_4bit_quant_type": "nf4",
-                    "bnb_4bit_compute_dtype": torch.bfloat16,
+                    "bnb_4bit_compute_dtype": dtype,
                 }
             bnb_config = BitsAndBytesConfig(**quantization_config)
 
-        try:
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                self.model_id
-            )
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.model_id
+        ).to(self.device)
 
-            if quantize:
-                self.model = AutoModelForCausalLM.from_pretrained(
-                    self.model_id,
-                    quantization_config=bnb_config,
-                    *args,
-                    **kwargs,
-                ).to(self.device)
-            else:
-                self.model = AutoModelForCausalLM.from_pretrained(
-                    self.model_id, *args, **kwargs
-                ).to(self.device)
+        if quantize:
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_id,
+                quantization_config=bnb_config,
+                *args,
+                **kwargs,
+            ).to(self.device)
+        else:
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_id, *args, **kwargs
+            ).to(self.device)
 
-        except Exception as e:
-            # self.logger.error(f"Failed to load the model or the tokenizer: {e}")
-            # raise
-            print(
-                colored(
-                    (
-                        "Failed to load the model and or the"
-                        f" tokenizer: {e}"
-                    ),
-                    "red",
-                )
-            )
 
     def print_error(self, error: str):
         """Print error"""
@@ -276,8 +264,7 @@ class HuggingfaceLLM(AbstractLLM):
                         *args,
                         **kwargs,
                     )
-
-            del inputs
+                    
             return self.tokenizer.decode(
                 outputs[0], skip_special_tokens=True
             )
diff --git a/swarms/structs/block_wrapper.py b/swarms/structs/block_wrapper.py
index cae44edf..a6811b8f 100644
--- a/swarms/structs/block_wrapper.py
+++ b/swarms/structs/block_wrapper.py
@@ -1,5 +1,6 @@
-from typing import Callable, Any
 import logging
+from functools import wraps
+from typing import Any, Callable
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -20,6 +21,7 @@ def block(
         Callable[..., Any]: The transformed function.
     """
 
+    @wraps(function)
     def wrapper(*args, **kwargs):
         # Here you can add code to execute the function on various hardwares
         # For now, we'll just call the function normally
diff --git a/swarms/structs/blockslist.py b/swarms/structs/blockslist.py
index 93ab0afa..8448454c 100644
--- a/swarms/structs/blockslist.py
+++ b/swarms/structs/blockslist.py
@@ -36,6 +36,15 @@ class BlocksList(BaseStructure):
         get_by_parent_name(parent_name: str): Get blocks by parent name.
         get_by_parent_type(parent_type: str): Get blocks by parent type.
         get_by_parent_description(parent_description: str): Get blocks by parent description.
+
+
+    Examples:
+    >>> from swarms.structs.block import Block
+    >>> from swarms.structs.blockslist import BlocksList
+    >>> block = Block("block", "A block")
+    >>> blockslist = BlocksList("blockslist", "A list of blocks", [block])
+    >>> blockslist
+
     """
 
     def __init__(
@@ -47,8 +56,10 @@ class BlocksList(BaseStructure):
         **kwargs,
     ):
         super().__init__(name=name, description=description, **kwargs)
-        self.parent = parent
+        self.name = name
+        self.description = description
         self.blocks = blocks
+        self.parent = parent
 
     def add(self, block: Any):
         self.blocks.append(block)
@@ -64,6 +75,26 @@ class BlocksList(BaseStructure):
 
     def get_all(self):
         return self.blocks
+    
+    def run_block(self, block: Any, task: str,  *args, **kwargs):
+        """Run the block for the specified task.
+
+        Args:
+            task (str): The task to be performed by the block.
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            The output of the block.
+
+        Raises:
+            Exception: If an error occurs during the execution of the block.
+        """
+        try:
+            return block.run(task, *args, **kwargs)
+        except Exception as error:
+            print(f"[Error] [Block] {error}")
+            raise error
 
     def get_by_name(self, name: str):
         return [block for block in self.blocks if block.name == name]
diff --git a/swarms/structs/model_parallizer.py b/swarms/structs/model_parallizer.py
index 3844f5b4..e115aa65 100644
--- a/swarms/structs/model_parallizer.py
+++ b/swarms/structs/model_parallizer.py
@@ -41,14 +41,12 @@ class ModelParallelizer:
     def __init__(
         self,
         llms: List[Callable] = None,
-        load_balancing: bool = False,
         retry_attempts: int = 3,
         iters: int = None,
         *args,
         **kwargs,
     ):
         self.llms = llms
-        self.load_balancing = load_balancing
         self.retry_attempts = retry_attempts
         self.iters = iters
         self.last_responses = None
@@ -151,16 +149,6 @@ class ModelParallelizer:
             )
         )
 
-    def enable_load_balancing(self):
-        """Enable load balancing among LLMs."""
-        self.load_balancing = True
-        logger.info("Load balancing enabled.")
-
-    def disable_load_balancing(self):
-        """Disable load balancing."""
-        self.load_balancing = False
-        logger.info("Load balancing disabled.")
-
     async def arun(self, task: str):
         """Asynchronous run the task string"""
         loop = asyncio.get_event_loop()
diff --git a/swarms/telemetry/auto_upgrade_swarms.py b/swarms/telemetry/auto_upgrade_swarms.py
index aead795b..210cf7c5 100644
--- a/swarms/telemetry/auto_upgrade_swarms.py
+++ b/swarms/telemetry/auto_upgrade_swarms.py
@@ -6,6 +6,10 @@ def auto_update():
     """auto update swarms"""
     try:
         if check_for_update():
+            print(
+                "There is a new version of swarms available!"
+                " Downloading..."
+            )
             subprocess.run(["pip", "install", "--upgrade", "swarms"])
     except Exception as e:
         print(e)
diff --git a/swarms/telemetry/bootup.py b/swarms/telemetry/bootup.py
new file mode 100644
index 00000000..edcd1aca
--- /dev/null
+++ b/swarms/telemetry/bootup.py
@@ -0,0 +1,8 @@
+from swarms.utils.disable_logging import disable_logging
+from swarms.telemetry.auto_upgrade_swarms import auto_update
+
+
+def bootup():
+    """Bootup swarms"""
+    disable_logging()
+    auto_update()
diff --git a/swarms/telemetry/check_update.py b/swarms/telemetry/check_update.py
index a9b6386e..2cdbd986 100644
--- a/swarms/telemetry/check_update.py
+++ b/swarms/telemetry/check_update.py
@@ -40,7 +40,3 @@ def check_for_update():
     return version.parse(latest_version) > version.parse(
         current_version
     )
-
-
-# out = check_for_update()
-# print(out)

From d0b09ea02927c7256599faa0152df48befa69c3c Mon Sep 17 00:00:00 2001
From: Eternal Reclaimer <98760976+kyegomez@users.noreply.github.com>
Date: Fri, 5 Jan 2024 17:37:39 -0500
Subject: [PATCH 7/8] Update team.py

---
 swarms/structs/team.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/swarms/structs/team.py b/swarms/structs/team.py
index 36c773e2..3b71165d 100644
--- a/swarms/structs/team.py
+++ b/swarms/structs/team.py
@@ -1,7 +1,7 @@
 import json
 from typing import List, Optional
 
-from pydantic.v1 import BaseModel, Field, Json, root_validator
+from pydantic import BaseModel, Field, Json, root_validator
 
 from swarms.structs.agent import Agent
 from swarms.structs.task import Task

From 91a36c855702df41516f8b2cee0ce74bc7f8c7c4 Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Fri, 5 Jan 2024 18:00:34 -0500
Subject: [PATCH 8/8] [FIX][HuggingfaceLLM]

---
 block.py                     | 2 +-
 pyproject.toml               | 2 +-
 swarms/models/huggingface.py | 5 ++---
 swarms/structs/blockslist.py | 4 ++--
 4 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/block.py b/block.py
index c5a76910..93894245 100644
--- a/block.py
+++ b/block.py
@@ -94,4 +94,4 @@ blocks_by_parent_description = swarm.get_by_parent_description(
 
 # Run the block in the swarm
 inference = swarm.run_block(toolagent, "Hello World")
-print(inference)
\ No newline at end of file
+print(inference)
diff --git a/pyproject.toml b/pyproject.toml
index e0602d43..fb6cc323 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms"
-version = "3.4.1"
+version = "3.4.2"
 description = "Swarms - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <kye@apac.ai>"]
diff --git a/swarms/models/huggingface.py b/swarms/models/huggingface.py
index 8e74cff5..cdfa9de3 100644
--- a/swarms/models/huggingface.py
+++ b/swarms/models/huggingface.py
@@ -131,7 +131,7 @@ class HuggingfaceLLM(AbstractLLM):
         temperature: float = 0.7,
         top_k: int = 40,
         top_p: float = 0.8,
-        dtype = torch.bfloat16,
+        dtype=torch.bfloat16,
         *args,
         **kwargs,
     ):
@@ -189,7 +189,6 @@ class HuggingfaceLLM(AbstractLLM):
                 self.model_id, *args, **kwargs
             ).to(self.device)
 
-
     def print_error(self, error: str):
         """Print error"""
         print(colored(f"Error: {error}", "red"))
@@ -264,7 +263,7 @@ class HuggingfaceLLM(AbstractLLM):
                         *args,
                         **kwargs,
                     )
-                    
+
             return self.tokenizer.decode(
                 outputs[0], skip_special_tokens=True
             )
diff --git a/swarms/structs/blockslist.py b/swarms/structs/blockslist.py
index 8448454c..b2a4db08 100644
--- a/swarms/structs/blockslist.py
+++ b/swarms/structs/blockslist.py
@@ -75,8 +75,8 @@ class BlocksList(BaseStructure):
 
     def get_all(self):
         return self.blocks
-    
-    def run_block(self, block: Any, task: str,  *args, **kwargs):
+
+    def run_block(self, block: Any, task: str, *args, **kwargs):
         """Run the block for the specified task.
 
         Args: