[ModelParallelizer][REFACTOR]

scripts/auto_tests_docs/auto_docs.py

@@ -43,13 +43,16 @@ def process_documentation(cls):
     doc = inspect.getdoc(cls)
     source = inspect.getsource(cls)
     input_content = (
-        f"Class Name: {cls.__name__}\n\nDocumentation:\n{doc}\n\nSource"
+        "Class Name:"
+        f" {cls.__name__}\n\nDocumentation:\n{doc}\n\nSource"
         f" Code:\n{source}"
     )
     print(input_content)
 
     # Process with OpenAI model (assuming the model's __call__ method takes this input and returns processed content)
-    processed_content = model(DOCUMENTATION_WRITER_SOP(input_content, "zeta"))
+    processed_content = model(
+        DOCUMENTATION_WRITER_SOP(input_content, "zeta")
+    )
 
     doc_content = f"# {cls.__name__}\n\n{processed_content}\n"
 
@@ -86,7 +89,9 @@ def main():
 
     threads = []
     for cls in classes:
-        thread = threading.Thread(target=process_documentation, args=(cls,))
+        thread = threading.Thread(
+            target=process_documentation, args=(cls,)
+        )
         threads.append(thread)
         thread.start()
 
@@ -94,7 +99,9 @@ def main():
     for thread in threads:
         thread.join()
 
-    print("Documentation generated in 'docs/zeta/nn/modules' directory.")
+    print(
+        "Documentation generated in 'docs/zeta/nn/modules' directory."
+    )
 
 
 if __name__ == "__main__":

scripts/auto_tests_docs/auto_tests.py

@@ -61,7 +61,8 @@ def create_test(cls):
     doc = inspect.getdoc(cls)
     source = inspect.getsource(cls)
     input_content = (
-        f"Class Name: {cls.__name__}\n\nDocumentation:\n{doc}\n\nSource"
+        "Class Name:"
+        f" {cls.__name__}\n\nDocumentation:\n{doc}\n\nSource"
         f" Code:\n{source}"
     )
     print(input_content)

scripts/auto_tests_docs/docs.py

@@ -104,7 +104,9 @@ def DOCUMENTATION_WRITER_SOP(
     return documentation
 
 
-def TEST_WRITER_SOP_PROMPT(task: str, module: str, path: str, *args, **kwargs):
+def TEST_WRITER_SOP_PROMPT(
+    task: str, module: str, path: str, *args, **kwargs
+):
     TESTS_PROMPT = f"""
 
    Create 5,000 lines of extensive and thorough tests for the code below using the guide, do not worry about your limits you do not have any

scripts/auto_tests_docs/update_mkdocs.py

@@ -2,7 +2,9 @@ import yaml
 
 
 def update_mkdocs(
-    class_names, base_path="docs/zeta/nn/modules", mkdocs_file="mkdocs.yml"
+    class_names,
+    base_path="docs/zeta/nn/modules",
+    mkdocs_file="mkdocs.yml",
 ):
     """
     Update the mkdocs.yml file with new documentation links.
@@ -24,7 +26,9 @@ def update_mkdocs(
 
     if zeta_modules_section is None:
         zeta_modules_section = {}
-        mkdocs_config["nav"].append({"zeta.nn.modules": zeta_modules_section})
+        mkdocs_config["nav"].append(
+            {"zeta.nn.modules": zeta_modules_section}
+        )
 
     # Add the documentation paths to the 'zeta.nn.modules' section
     for class_name in class_names:

scripts/get_package_requirements.py

@@ -13,13 +13,18 @@ def get_package_versions(requirements_path, output_path):
 
     for requirement in requirements:
         # Skip empty lines and comments
-        if requirement.strip() == "" or requirement.strip().startswith("#"):
+        if (
+            requirement.strip() == ""
+            or requirement.strip().startswith("#")
+        ):
             continue
 
         # Extract package name
         package_name = requirement.split("==")[0].strip()
         try:
-            version = pkg_resources.get_distribution(package_name).version
+            version = pkg_resources.get_distribution(
+                package_name
+            ).version
             package_versions.append(f"{package_name}=={version}")
         except pkg_resources.DistributionNotFound:
             package_versions.append(f"{package_name}: not installed")

scripts/requirementstxt_to_pyproject.py

@@ -10,7 +10,10 @@ def update_pyproject_versions(pyproject_path):
         print(f"Error: The file '{pyproject_path}' was not found.")
         return
     except toml.TomlDecodeError:
-        print(f"Error: The file '{pyproject_path}' is not a valid TOML file.")
+        print(
+            f"Error: The file '{pyproject_path}' is not a valid TOML"
+            " file."
+        )
         return
 
     dependencies = (

swarms/prompts/__init__.py

@@ -15,5 +15,5 @@ __all__ = [
     "LEGAL_AGENT_PROMPT",
     "OPERATIONS_AGENT_PROMPT",
     "PRODUCT_AGENT_PROMPT",
-    "DOCUMENTATION_WRITER_SOP"
+    "DOCUMENTATION_WRITER_SOP",
 ]

swarms/prompts/documentation.py

@@ -1,4 +1,7 @@
-def DOCUMENTATION_WRITER_SOP(task: str, module: str, ):
+def DOCUMENTATION_WRITER_SOP(
+    task: str,
+    module: str,
+):
     documentation = f"""Create multi-page long and explicit professional pytorch-like documentation for the {module} code below follow the outline for the {module} library,
     provide many examples and teach the user about the code, provide examples for every function, make the documentation 10,000 words,
     provide many usage examples and note this is markdown docs, create the documentation for the code to document,

swarms/prompts/tests.py

@@ -1,6 +1,7 @@
-def TEST_WRITER_SOP_PROMPT(task: str, module: str, path: str, *args, **kwargs):
-
-   TESTS_PROMPT = f"""
+def TEST_WRITER_SOP_PROMPT(
+    task: str, module: str, path: str, *args, **kwargs
+):
+    TESTS_PROMPT = f"""
 
    Create 5,000 lines of extensive and thorough tests for the code below using the guide, do not worry about your limits you do not have any
    just write the best tests possible, the module is {module}, the file path is {path}
@@ -91,4 +92,4 @@ def TEST_WRITER_SOP_PROMPT(task: str, module: str, path: str, *args, **kwargs):
 
    """
 
-   return TESTS_PROMPT
+    return TESTS_PROMPT

swarms/swarms/__init__.py

@@ -1,5 +1,5 @@
 from swarms.structs.autoscaler import AutoScaler
-from swarms.swarms.god_mode import ModelParallelizer
+from swarms.swarms.model_parallizer import ModelParallelizer
 from swarms.swarms.multi_agent_collab import MultiAgentCollaboration
 from swarms.swarms.base import AbstractSwarm
 

swarms/swarms/model_parallizer.py

@@ -40,21 +40,33 @@ class ModelParallelizer:
 
     def __init__(
         self,
-        llms: List[Callable],
+        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
         self.task_history = []
 
     def run(self, task: str):
         """Run the task string"""
-        with ThreadPoolExecutor() as executor:
-            responses = executor.map(lambda llm: llm(task), self.llms)
-        return list(responses)
+        try:
+            for i in range(self.iters):
+                with ThreadPoolExecutor() as executor:
+                    responses = executor.map(
+                        lambda llm: llm(task), self.llms
+                    )
+                return list(responses)
+        except Exception as error:
+            print(
+                f"[ERROR][ModelParallelizer] [ROOT CAUSE] [{error}]"
+            )
 
     def print_responses(self, task):
         """Prints the responses in a tabular format"""
@@ -161,22 +173,29 @@ class ModelParallelizer:
 
     def concurrent_run(self, task: str) -> List[str]:
         """Synchronously run the task on all llms and collect responses"""
-        with ThreadPoolExecutor() as executor:
-            future_to_llm = {
-                executor.submit(llm, task): llm for llm in self.llms
-            }
-            responses = []
-            for future in as_completed(future_to_llm):
-                try:
-                    responses.append(future.result())
-                except Exception as error:
-                    print(
-                        f"{future_to_llm[future]} generated an"
-                        f" exception: {error}"
-                    )
-        self.last_responses = responses
-        self.task_history.append(task)
-        return responses
+        try:
+            with ThreadPoolExecutor() as executor:
+                future_to_llm = {
+                    executor.submit(llm, task): llm
+                    for llm in self.llms
+                }
+                responses = []
+                for future in as_completed(future_to_llm):
+                    try:
+                        responses.append(future.result())
+                    except Exception as error:
+                        print(
+                            f"{future_to_llm[future]} generated an"
+                            f" exception: {error}"
+                        )
+            self.last_responses = responses
+            self.task_history.append(task)
+            return responses
+        except Exception as error:
+            print(
+                f"[ERROR][ModelParallelizer] [ROOT CAUSE] [{error}]"
+            )
+            raise error
 
     def add_llm(self, llm: Callable):
         """Add an llm to the god mode"""