diff --git a/docs/swarms/models/mistral.md b/docs/swarms/models/mistral.md deleted file mode 100644 index 6c38cd97..00000000 --- a/docs/swarms/models/mistral.md +++ /dev/null @@ -1,284 +0,0 @@ -# `Mistral` Documentation - -## Table of Contents - -1. [Introduction](#introduction) -2. [Overview](#overview) -3. [Class Definition](#class-definition) - - [Mistral Class](#mistral-class) - - [Initialization Parameters](#initialization-parameters) -4. [Functionality and Usage](#functionality-and-usage) - - [Loading the Model](#loading-the-model) - - [Running the Model](#running-the-model) - - [Chatting with the Agent](#chatting-with-the-agent) -5. [Additional Information](#additional-information) -6. [Examples](#examples) - - [Example 1: Initializing Mistral](#example-1-initializing-mistral) - - [Example 2: Running a Task](#example-2-running-a-task) - - [Example 3: Chatting with the Agent](#example-3-chatting-with-the-agent) -7. [References and Resources](#references-and-resources) - ---- - -## 1. Introduction - -Welcome to the documentation for Mistral, a powerful language model-based AI agent. Mistral leverages the capabilities of large language models to generate text-based responses to queries and tasks. This documentation provides a comprehensive guide to understanding and using the Mistral AI agent. - -### 1.1 Purpose - -Mistral is designed to assist users by generating coherent and contextually relevant text based on user inputs or tasks. It can be used for various natural language understanding and generation tasks, such as chatbots, text completion, question answering, and content generation. - -### 1.2 Key Features - -- Utilizes large pre-trained language models. -- Supports GPU acceleration for faster processing. -- Provides an easy-to-use interface for running tasks and engaging in chat-based conversations. -- Offers fine-grained control over response generation through temperature and maximum length settings. - ---- - -## 2. Overview - -Before diving into the details of the Mistral AI agent, let's provide an overview of its purpose and functionality. - -Mistral is built on top of powerful language models, such as GPT-3. It allows you to: - -- Generate text-based responses to tasks and queries. -- Control the temperature of response generation for creativity. -- Set a maximum length for generated responses. -- Engage in chat-based conversations with the AI agent. -- Utilize GPU acceleration for faster inference. - -In the following sections, we will explore the class definition, its initialization parameters, and how to use Mistral effectively. - ---- - -## 3. Class Definition - -Mistral consists of a single class, the `Mistral` class. This class provides methods for initializing the agent, loading the pre-trained model, and running tasks. - -### 3.1 Mistral Class - -```python -class Mistral: - """ - Mistral - - model = Mistral(device="cuda", use_flash_attention=True, temperature=0.7, max_length=200) - task = "My favorite condiment is" - result = model.run(task) - print(result) - """ - - def __init__( - self, - ai_name: str = "Node Model Agent", - system_prompt: str = None, - model_name: str = "mistralai/Mistral-7B-v0.1", - device: str = "cuda", - use_flash_attention: bool = False, - temperature: float = 1.0, - max_length: int = 100, - do_sample: bool = True, - ): - """ - Initializes the Mistral AI agent. - - Parameters: - - ai_name (str): The name or identifier of the AI agent. Default: "Node Model Agent". - - system_prompt (str): A system-level prompt for context (e.g., conversation history). Default: None. - - model_name (str): The name of the pre-trained language model to use. Default: "mistralai/Mistral-7B-v0.1". - - device (str): The device for model inference, such as "cuda" or "cpu". Default: "cuda". - - use_flash_attention (bool): If True, enables flash attention for faster inference. Default: False. - - temperature (float): A value controlling the creativity of generated text. Default: 1.0. - - max_length (int): The maximum length of generated text. Default: 100. - - do_sample (bool): If True, uses sampling for text generation. Default: True. - """ -``` - -### 3.2 Initialization Parameters - -- `ai_name` (str): The name or identifier of the AI agent. This name can be used to distinguish between different agents if multiple instances are used. The default value is "Node Model Agent". - -- `system_prompt` (str): A system-level prompt that provides context for the AI agent. This can be useful for maintaining a conversation history or providing context for the current task. By default, it is set to `None`. - -- `model_name` (str): The name of the pre-trained language model to use. The default value is "mistralai/Mistral-7B-v0.1", which points to a specific version of the Mistral model. - -- `device` (str): The device on which the model should perform inference. You can specify "cuda" for GPU acceleration or "cpu" for CPU-only inference. The default is "cuda", assuming GPU availability. - -- `use_flash_attention` (bool): If set to `True`, Mistral uses flash attention for faster inference. This is beneficial when low-latency responses are required. The default is `False`. - -- `temperature` (float): The temperature parameter controls the creativity of the generated text. Higher values (e.g., 1.0) produce more random output, while lower values (e.g., 0.7) make the output more focused and deterministic. The default value is 1.0. - -- `max_length` (int): This parameter sets the maximum length of the generated text. It helps control the length of responses. The default value is 100. - -- `do_sample` (bool): If set to `True`, Mistral uses sampling during text generation. Sampling introduces randomness into the generated text. The default is `True`. - ---- - -## 4. Functionality and Usage - -Now that we've introduced the Mistral class and its parameters, let's explore how to use Mistral for various tasks. - -### 4.1 Loading the Model - -The `Mistral` class handles the loading of the pre-trained language model during initialization. You do not need to explicitly load the model. Simply create an instance of `Mistral`, and it will take care of loading the model into memory. - -### 4.2 Running the Model - -Mistral provides two methods for running the model: - -#### 4.2.1 `run` Method - -The `run` method is used to generate text-based responses to a given task or input. It takes a single string parameter, `task`, and returns the generated text as a string. - -```python -def run(self, task: str) -> str: - """ - Run the model on a given task. - - Parameters: - - task (str): The task or query for which to generate a response. - - Returns: - - str: The generated text response. - """ -``` - -Example: - -```python -from swarms.models import Mistral - -model = Mistral() -task = "Translate the following English text to French: 'Hello, how are you?'" -result = model.run(task) -print(result) -``` - -#### 4.2.2 `__call__` Method - -The `__call__` method provides a more concise way to run the model on a given task. You can use it by simply calling the `Mistral` instance with a task string. - -Example: - -```python -model = Mistral() -task = "Generate a summary of the latest research paper on AI ethics." -result = model(task) -print(result) -``` - -### 4.3 Chatting with the Agent - -Mistral supports chat-based interactions with the AI agent. You can send a series of messages to the agent, and it will respond accordingly. The `chat` method handles these interactions. - -#### `chat` Method - -The `chat` method allows you to engage in chat-based conversations with the AI agent. You can send messages to the agent, and it will respond with text-based messages. - -```python -def chat(self, msg: str = None, streaming: bool = False) -> str: - """ - Run a chat conversation with the agent. - - Parameters: - - msg (str, optional): The message to send to the agent. Defaults to None. - - streaming (bool, optional): Whether to stream the response token by token. Defaults to False. - - Returns: - - str: The response from the agent. - """ -``` - -Example: - -```python -model = Mistral() -conversation = [ - "Tell me a joke.", - "What's the weather like today?", - "Translate 'apple' to Spanish.", -] -for user_message in conversation: - response = model.chat(user_message) - print(f"User: {user_message}") - print(f"Agent: {response}") -``` - ---- - -## 5. Additional Information - -Here are some additional tips and information for using Mistral effectively: - -- Mistral uses a specific pre-trained model ("mistralai/Mistral-7B-v0.1" by default). You can explore other available models and choose one that best suits your task. - -- The `temperature` parameter controls the randomness of generated text. Experiment with different values to achieve the desired level of creativity in responses. - -- Be cautious with `max_length`, especially if you set it to a very high value, as it may lead to excessively long responses. - -- Ensure that you have the required libraries, such as `torch` and `transformers`, installed to use Mistral successfully. - -- Consider providing a system-level prompt when engaging in chat-based conversations to provide context for the agent. - ---- - -## 6. Examples - -In this section, we provide practical examples to illustrate how to use Mistral for various tasks. - -### 6.1 Example 1: Initializing Mistral - -In this example, we initialize the Mistral AI agent with custom settings: - -```python -from swarms.models import Mistral - -model = Mistral( - ai_name="My AI Assistant", - device="cpu", - temperature=0.8, - max_length=150, -) -``` - -### 6.2 Example 2: Running a Task - -Here, we run a text generation task using Mistral: - -```python -model = Mistral() -task = "Summarize the main findings of the recent climate change report." -result = model.run(task) -print(result) -``` - -### 6.3 Example 3: Chatting with the Agent - -Engage in a chat-based conversation with Mistral: - -```python -model = Mistral() -conversation = [ - "Tell me a joke.", - "What's the latest news?", - "Translate 'cat' to French.", -] -for user_message in conversation: - response = model.chat(user_message) - print(f"User: {user_message}") - print(f"Agent: {response}") -``` - ---- - -## 7. References and Resources - -Here are some references and resources for further information on Mistral and related topics: - -- [Mistral GitHub Repository](https://github.com/mistralai/mistral): Official Mistral repository for updates and contributions. -- [Hugging Face Transformers](https://huggingface.co/transformers/): Documentation and models for various transformers, including Mistral's parent models. -- [PyTorch Official Website](https://pytorch.org/): Official website for PyTorch, the deep learning framework used in Mistral. - -This concludes the documentation for the Mistral AI agent. You now have a comprehensive understanding of how to use Mistral for text generation and chat-based interactions. If you have any further questions or need assistance, please refer to the provided references and resources. Happy AI modeling! \ No newline at end of file diff --git a/docs/swarms/models/mixtral.md b/docs/swarms/models/mixtral.md deleted file mode 100644 index aa1b64d3..00000000 --- a/docs/swarms/models/mixtral.md +++ /dev/null @@ -1,76 +0,0 @@ -# Module Name: Mixtral - -## Introduction -The Mixtral module is a powerful language model designed for text generation tasks. It leverages the MistralAI Mixtral-8x7B pre-trained model to generate high-quality text based on user-defined tasks or prompts. In this documentation, we will provide a comprehensive overview of the Mixtral module, including its architecture, purpose, arguments, and detailed usage examples. - -## Purpose -The Mixtral module is designed to facilitate text generation tasks using state-of-the-art language models. Whether you need to generate creative content, draft text for various applications, or simply explore the capabilities of Mixtral, this module serves as a versatile and efficient solution. With its easy-to-use interface, you can quickly generate text for a wide range of applications. - -## Architecture -The Mixtral module is built on top of the MistralAI Mixtral-8x7B pre-trained model. It utilizes a deep neural network architecture with 8 layers and 7 attention heads to generate coherent and contextually relevant text. The model is capable of handling a variety of text generation tasks, from simple prompts to more complex content generation. - -## Class Definition -### `Mixtral(model_name: str = "mistralai/Mixtral-8x7B-v0.1", max_new_tokens: int = 500)` - -#### Parameters -- `model_name` (str, optional): The name or path of the pre-trained Mixtral model. Default is "mistralai/Mixtral-8x7B-v0.1". -- `max_new_tokens` (int, optional): The maximum number of new tokens to generate. Default is 500. - -## Functionality and Usage -The Mixtral module offers a straightforward interface for text generation. It accepts a task or prompt as input and returns generated text based on the provided input. - -### `run(task: Optional[str] = None, **kwargs) -> str` - -#### Parameters -- `task` (str, optional): The task or prompt for text generation. - -#### Returns -- `str`: The generated text. - -## Usage Examples -### Example 1: Basic Usage - -```python -from swarms.models import Mixtral - -# Initialize the Mixtral model -mixtral = Mixtral() - -# Generate text for a simple task -generated_text = mixtral.run("Generate a creative story.") -print(generated_text) -``` - -### Example 2: Custom Model - -You can specify a custom pre-trained model by providing the `model_name` parameter. - -```python -custom_model_name = "model_name" -mixtral_custom = Mixtral(model_name=custom_model_name) - -generated_text = mixtral_custom.run("Generate text with a custom model.") -print(generated_text) -``` - -### Example 3: Controlling Output Length - -You can control the length of the generated text by adjusting the `max_new_tokens` parameter. - -```python -mixtral_length = Mixtral(max_new_tokens=100) - -generated_text = mixtral_length.run("Generate a short text.") -print(generated_text) -``` - -## Additional Information and Tips -- It's recommended to use a descriptive task or prompt to guide the text generation process. -- Experiment with different prompt styles and lengths to achieve the desired output. -- You can fine-tune Mixtral on specific tasks if needed, although pre-trained models often work well out of the box. -- Monitor the `max_new_tokens` parameter to control the length of the generated text. - -## Conclusion -The Mixtral module is a versatile tool for text generation tasks, powered by the MistralAI Mixtral-8x7B pre-trained model. Whether you need creative writing, content generation, or assistance with text-based tasks, Mixtral can help you achieve your goals. With a simple interface and flexible parameters, it's a valuable addition to your text generation toolkit. - -If you encounter any issues or have questions about using Mixtral, please refer to the MistralAI documentation or reach out to their support team for further assistance. Happy text generation with Mixtral! \ No newline at end of file diff --git a/docs/swarms/models/mpt.md b/docs/swarms/models/mpt.md deleted file mode 100644 index 0592284b..00000000 --- a/docs/swarms/models/mpt.md +++ /dev/null @@ -1,120 +0,0 @@ -# `MPT7B` -============================================== - -The `MPT7B` class is a powerful tool for generating text using pre-trained models. It leverages the `transformers` library from Hugging Face to load models and tokenizers, and to perform the text generation. The class is designed to be flexible and easy to use, with a variety of methods for generating text both synchronously and asynchronously. - -## Class Definition ----------------- - -``` -class MPT7B: - def __init__(self, model_name: str, tokenizer_name: str, max_tokens: int = 100) - def run(self, task: str, *args, **kwargs) -> str - async def run_async(self, task: str, *args, **kwargs) -> str - def generate(self, prompt: str) -> str - async def generate_async(self, prompt: str) -> str - def __call__(self, task: str, *args, **kwargs) -> str - async def __call_async__(self, task: str, *args, **kwargs) -> str - def batch_generate(self, prompts: list, temperature: float = 1.0) -> list - def unfreeze_model(self) -``` - - -## Class Parameters ----------------- - -| Parameter | Type | Description | -| --- | --- | --- | -| `model_name` | str | Name of the pre-trained model to use. | -| `tokenizer_name` | str | Name of the tokenizer to use. | -| `max_tokens` | int | Maximum number of tokens to generate. Default is 100. | - -## Class Methods -------------- - -| Method | Returns | Description | -| --- | --- | --- | -| `run(task: str, *args, **kwargs)` | str | Run the model with the specified task and arguments. | -| `run_async(task: str, *args, **kwargs)` | str | Run the model asynchronously with the specified task and arguments. | -| `generate(prompt: str)` | str | Generate text from the given prompt. | -| `generate_async(prompt: str)` | str | Generate text asynchronously from the given prompt. | -| `__call__(task: str, *args, **kwargs)` | str | Call the model with the specified task and arguments. | -| `__call_async__(task: str, *args, **kwargs)` | str | Call the model asynchronously with the specified task and arguments. | -| `batch_generate(prompts: list, temperature: float = 1.0)` | list | Generate text for a batch of prompts. | -| `unfreeze_model()` | None | Unfreeze the model for fine-tuning. | - -## Usage Examples --------------- - -### Example 1: Basic Text Generation - -```python -from swarms.models import MPT7B - -# Initialize the MPT7B class -mpt = MPT7B("mosaicml/mpt-7b-storywriter", "EleutherAI/gpt-neox-20b", max_tokens=150) - -# Generate text -output = mpt.run("generate", "Once upon a time in a land far, far away...") -print(output) -``` - -### Example 2: Batch Text Generation - -```pyton -from swarms.models import MPT7B - -# Initialize the MPT7B class -mpt = MPT7B('mosaicml/mpt-7b-storywriter', 'EleutherAI/gpt-neox-20b', max_tokens=150) - -# Generate text for a batch of prompts -prompts = ['In the deep jungles,', 'At the heart of the city,'] -outputs = mpt.batch_generate(prompts, temperature=0.7) -print(outputs) -``` - -### Example 3: Asynchronous Text Generation - -```python -import asyncio - -from swarms.models import MPT7B - -# Initialize the MPT7B class -mpt = MPT7B("mosaicml/mpt-7b-storywriter", "EleutherAI/gpt-neox-20b", max_tokens=150) - -# Generate text asynchronously -output = asyncio.run( - mpt.run_async("generate", "Once upon a time in a land far, far away...") -) -print(output) -``` - -## Additional Information ----------------------------------- - -The `batch_generate` method allows for generating text for multiple prompts at once. This can be more efficient than generating text for each prompt individually, especially when working with a large number of prompts. - -The `unfreeze_model` method is used to unfreeze the model for fine-tuning. By default, the model parameters are frozen to prevent them from being updated during training. Unfreezing the model allows the parameters to be updated, which can be useful for fine-tuning the model on a specific task. - -The `__call__` and `__call_async__` methods are convenience methods that allow the class instance to be called like a function. They simply call the `run` and `run_async` methods, respectively. - -## Architecture and Working ------------------------- - -The `MPT7B` class is designed to be a simple and flexible interface for text generation with pre-trained models. It encapsulates the complexity of loading models and tokenizers, setting up the text generation pipeline, and generating text. - -The class uses the `AutoModelForCausalLM` and `AutoTokenizer` classes from the `transformers` library to load the pre-trained model and tokenizer. The `pipeline` function is used to create a text generation pipeline with the loaded model and tokenizer. This pipeline is used to generate text from prompts. - -The `run` and `run_async` methods are the main entry points for using the class. They accept a task name and arbitrary arguments, and call the appropriate method based on the task name. The `generate` and `generate_async` methods perform the actual text generation. - -The `batch_generate` method allows for generating text for multiple prompts at once. This can be more efficient than generating text for each prompt individually, especially when working with a large number of prompts. - -The `unfreeze_model` method is used to unfreeze the model for fine-tuning. By default, the model parameters are frozen to prevent them from being updated during training. Unfreezing the model allows the parameters to be updated, which can be useful for fine-tuning the model on a specific task. - -The `__call__` and `__call_async__` methods are convenience methods that allow the class instance to be called like a function. They simply call the `run` and `run_async` methods, respectively. - -## Conclusion ----------- - -The `MPT7B` class provides a powerful and flexible interface for text generation with pre-trained models. It encapsulates the complexity of loading models and tokenizers, setting up the text generation pipeline, and generating text, making it easy to generate high-quality text with just a few lines of code. Whether you're generating text for a single prompt, a batch of prompts, or fine-tuning the model on a specific task, the `MPT7B` class has you covered. \ No newline at end of file diff --git a/docs/swarms/structs/json.md b/docs/swarms/structs/json.md deleted file mode 100644 index a495a0ce..00000000 --- a/docs/swarms/structs/json.md +++ /dev/null @@ -1,134 +0,0 @@ -# **Documentation for `swarms.structs.JSON` Class** - -The `swarms.structs.JSON` class is a helper class that provides a templated framework for creating new classes that deal with JSON objects and need to validate these objects against a JSON Schema. Being an abstract base class (ABC), the `JSON` class allows for the creation of subclasses that implement specific behavior while ensuring that they all adhere to a common interface, particularly the `validate` method. - -Given that documenting the entire code provided in full detail would exceed our platform's limitations, below is a generated documentation for the `JSON` class following the steps you provided. This is an outline and would need to be expanded upon to reach the desired word count and thoroughness in a full, professional documentation. - ---- - -## Introduction - -JSON (JavaScript Object Notation) is a lightweight data interchange format that is easy for humans to read and write and easy for machines to parse and generate. `swarms.structs.JSON` class aims to provide a basic structure for utilizing JSON and validating it against a pre-defined schema. This is essential for applications where data integrity and structure are crucial, such as configurations for applications, communications over networks, and data storage. - -## Class Definition - -| Parameter | Type | Description | -|---------------|------------|------------------------------------| -| `schema_path` | `str` | The path to the JSON schema file. | - -### `JSON.__init__(self, schema_path)` -Class constructor that initializes a `JSON` object with the specified JSON schema path. -```python -def __init__(self, schema_path): - self.schema_path = schema_path - self.schema = self.load_schema() -``` - -### `JSON.load_schema(self)` -Private method that loads and returns the JSON schema from the file specified at the `schema_path`. - -### `JSON.validate(self, data)` -Abstract method that needs to be implemented by subclasses to validate input `data` against the JSON schema. - ---- - -## Functionality and Usage - -### Why use `JSON` Class - -The `JSON` class abstracts away the details of loading and validating JSON data, allowing for easy implementation in any subclass that needs to handle JSON input. It sets up a standard for all subclasses to follow, ensuring consistency across different parts of code or different projects. - -By enforcing a JSON schema, the `JSON` class helps maintain the integrity of the data, catching errors early in the process of reading and writing JSON. - -### Step-by-step Guide - -1. Subclass the `JSON` class. -2. Provide an implementation for the `validate` method. -3. Use the provided schema to enforce required fields and types within your JSON data. - ---- - -## Example Usage - -### Implementing a Subclass - -Suppose we have a JSON Schema in `config_schema.json` for application configuration. - -```json -{ - "$schema": "http://json-schema.org/draft-07/schema#", - "type": "object", - "properties": { - "debug": { - "type": "boolean" - }, - "window_size": { - "type": "array", - "items": { - "type": "number" - }, - "minItems": 2, - "maxItems": 2 - } - }, - "required": ["debug", "window_size"] -} -``` - -Now we'll create a subclass `AppConfig` that uses this schema. - -```python -from swarms.structs import JSON - - -class AppConfig(JSON): - def __init__(self, schema_path): - super().__init__(schema_path) - - def validate(self, config_data): - # Here we'll use a JSON Schema validation library like jsonschema - from jsonschema import ValidationError, validate - - try: - validate(instance=config_data, schema=self.schema) - except ValidationError as e: - print(f"Invalid configuration: {e}") - return False - return True - - -# Main Example Usage - -if __name__ == "__main__": - config = {"debug": True, "window_size": [800, 600]} - - app_config = AppConfig("config_schema.json") - - if app_config.validate(config): - print("Config is valid!") - else: - print("Config is invalid.") -``` - -In this example, an `AppConfig` class that inherits from `JSON` is created. The `validate` method is implemented to check whether a configuration dictionary is valid against the provided schema. - -### Note - -- Validate real JSON data using this class in a production environment. -- Catch and handle any exceptions as necessary to avoid application crashes. -- Extend functionality within subclasses as required for your application. - ---- - -## Additional Information and Tips - -- Use detailed JSON Schemas for complex data validation. -- Use the jsonschema library for advanced validation features. - -## References and Resources - -- Official Python Documentation for ABCs: https://docs.python.org/3/library/abc.html -- JSON Schema: https://json-schema.org/ -- jsonschema Python package: https://pypi.org/project/jsonschema/ - -This generated documentation serves as a template and starting point intended for creating in-depth, practical documentation. Expanding upon each section, in practice, would involve deeper code examples, common patterns and pitfalls, and more thorough explanations of the `JSON` class internals and how to best utilize them in various real-world scenarios. diff --git a/tests/structs/test_majorityvoting.py b/tests/structs/test_majorityvoting.py deleted file mode 100644 index e69de29b..00000000 diff --git a/tests/structs/test_model_parallizer.py b/tests/structs/test_model_parallizer.py deleted file mode 100644 index a0840608..00000000 --- a/tests/structs/test_model_parallizer.py +++ /dev/null @@ -1,147 +0,0 @@ -import pytest - -from swarms.models import ( - GPT4VisionAPI, - HuggingfaceLLM, - Mixtral, - ZeroscopeTTV, -) -from swarms.structs.model_parallizer import ModelParallelizer - -# Initialize the models -custom_config = { - "quantize": True, - "quantization_config": {"load_in_4bit": True}, - "verbose": True, -} -huggingface_llm = HuggingfaceLLM( - model_id="NousResearch/Nous-Hermes-2-Vision-Alpha", - **custom_config, -) -mixtral = Mixtral(load_in_4bit=True, use_flash_attention_2=True) -gpt4_vision_api = GPT4VisionAPI(max_tokens=1000) -zeroscope_ttv = ZeroscopeTTV() - - -def test_init(): - mp = ModelParallelizer( - [ - huggingface_llm, - mixtral, - gpt4_vision_api, - zeroscope_ttv, - ] - ) - assert isinstance(mp, ModelParallelizer) - - -def test_run(): - mp = ModelParallelizer([huggingface_llm]) - result = mp.run( - "Create a list of known biggest risks of structural collapse" - " with references" - ) - assert isinstance(result, str) - - -def test_run_all(): - mp = ModelParallelizer( - [ - huggingface_llm, - mixtral, - gpt4_vision_api, - zeroscope_ttv, - ] - ) - result = mp.run_all( - "Create a list of known biggest risks of structural collapse" - " with references" - ) - assert isinstance(result, list) - assert len(result) == 5 - - -def test_add_llm(): - mp = ModelParallelizer([huggingface_llm]) - mp.add_llm(mixtral) - assert len(mp.llms) == 2 - - -def test_remove_llm(): - mp = ModelParallelizer([huggingface_llm, mixtral]) - mp.remove_llm(mixtral) - assert len(mp.llms) == 1 - - -def test_save_responses_to_file(tmp_path): - mp = ModelParallelizer([huggingface_llm]) - mp.run( - "Create a list of known biggest risks of structural collapse" - " with references" - ) - file = tmp_path / "responses.txt" - mp.save_responses_to_file(file) - assert file.read_text() != "" - - -def test_get_task_history(): - mp = ModelParallelizer([huggingface_llm]) - mp.run( - "Create a list of known biggest risks of structural collapse" - " with references" - ) - assert mp.get_task_history() == [ - "Create a list of known biggest risks of structural collapse" - " with references" - ] - - -def test_summary(capsys): - mp = ModelParallelizer([huggingface_llm]) - mp.run( - "Create a list of known biggest risks of structural collapse" - " with references" - ) - mp.summary() - captured = capsys.readouterr() - assert "Tasks History:" in captured.out - - -def test_enable_load_balancing(): - mp = ModelParallelizer([huggingface_llm]) - mp.enable_load_balancing() - assert mp.load_balancing is True - - -def test_disable_load_balancing(): - mp = ModelParallelizer([huggingface_llm]) - mp.disable_load_balancing() - assert mp.load_balancing is False - - -def test_concurrent_run(): - mp = ModelParallelizer([huggingface_llm, mixtral]) - result = mp.concurrent_run( - "Create a list of known biggest risks of structural collapse" - " with references" - ) - assert isinstance(result, list) - assert len(result) == 2 - - -def test_concurrent_run_no_task(): - mp = ModelParallelizer([huggingface_llm]) - with pytest.raises(TypeError): - mp.concurrent_run() - - -def test_concurrent_run_non_string_task(): - mp = ModelParallelizer([huggingface_llm]) - with pytest.raises(TypeError): - mp.concurrent_run(123) - - -def test_concurrent_run_empty_task(): - mp = ModelParallelizer([huggingface_llm]) - result = mp.concurrent_run("") - assert result == [""]