From 98e922f75cfb79783ec84694d87380a01c608351 Mon Sep 17 00:00:00 2001 From: Kye Date: Tue, 26 Mar 2024 10:38:48 -0700 Subject: [PATCH] [CLEANUP] --- swarms/models/zephyr.py | 107 --------------------------- swarms/structs/all_to_all_swarm.py | 0 swarms/structs/recursive_workflow.py | 1 + swarms/structs/sermon_swarm.py | 1 + swarms/structs/task_queue_base.py | 24 +++--- 5 files changed, 14 insertions(+), 119 deletions(-) delete mode 100644 swarms/models/zephyr.py delete mode 100644 swarms/structs/all_to_all_swarm.py diff --git a/swarms/models/zephyr.py b/swarms/models/zephyr.py deleted file mode 100644 index 205ec2e5..00000000 --- a/swarms/models/zephyr.py +++ /dev/null @@ -1,107 +0,0 @@ -"""Zephyr by HF""" - -import torch -from transformers import pipeline - - -class Zephyr: - """ - Zehpyr model from HF - - - Args: - max_new_tokens(int) = Number of max new tokens - temperature(float) = temperature of the LLM - top_k(float) = top k of the model set to 50 - top_p(float) = top_p of the model set to 0.95 - - - - Usage: - >>> model = Zephyr() - >>> output = model("Generate hello world in python") - - - """ - - def __init__( - self, - model_name: str = "HuggingFaceH4/zephyr-7b-alpha", - tokenize: bool = False, - add_generation_prompt: bool = True, - system_prompt: str = "You are a friendly chatbot who always responds in the style of a pirate", - max_new_tokens: int = 300, - temperature: float = 0.5, - top_k: float = 50, - top_p: float = 0.95, - do_sample: bool = True, - *args, - **kwargs, - ): - super().__init__() - self.model_name = model_name - self.tokenize = tokenize - self.add_generation_prompt = add_generation_prompt - self.system_prompt = system_prompt - self.max_new_tokens = max_new_tokens - self.temperature = temperature - self.top_k = top_k - self.top_p = top_p - self.do_sample = do_sample - - self.pipe = pipeline( - "text-generation", - model=self.model_name, - torch_dtype=torch.bfloat16, - device_map="auto", - ) - self.messages = [ - { - "role": "system", - "content": f"{self.system_prompt}\n\nUser:", - }, - ] - - def __call__(self, task: str): - """Call the model""" - prompt = self.pipe.tokenizer.apply_chat_template( - self.messages, - tokenize=self.tokenize, - add_generation_prompt=self.add_generation_prompt, - ) - outputs = self.pipe( - prompt - ) # max_new_token=self.max_new_tokens) - print(outputs[0]["generated_text"]) - - def chat(self, message: str): - """ - Adds a user message to the conversation and generates a chatbot response. - """ - # Add the user message to the conversation - self.messages.append({"role": "user", "content": message}) - - # Apply the chat template to format the messages - prompt = self.pipe.tokenizer.apply_chat_template( - self.messages, - tokenize=self.tokenize, - add_generation_prompt=self.add_generation_prompt, - ) - - # Generate a response - outputs = self.pipe( - prompt, - max_new_tokens=self.max_new_tokens, - do_sample=self.do_sample, - temperature=self.temperature, - top_k=self.top_k, - top_p=self.top_p, - ) - - # Extract the generated text - generated_text = outputs[0]["generated_text"] - - # Optionally, you could also add the chatbot's response to the messages list - # However, the below line should be adjusted to extract the chatbot's response only - # self.messages.append({"role": "bot", "content": generated_text}) - return generated_text diff --git a/swarms/structs/all_to_all_swarm.py b/swarms/structs/all_to_all_swarm.py deleted file mode 100644 index e69de29b..00000000 diff --git a/swarms/structs/recursive_workflow.py b/swarms/structs/recursive_workflow.py index 60d33fe4..ecb679bc 100644 --- a/swarms/structs/recursive_workflow.py +++ b/swarms/structs/recursive_workflow.py @@ -36,6 +36,7 @@ class RecursiveWorkflow(BaseStructure): stopping_conditions: callable = None, ): self.stop_token = stop_token + self.stopping_conditions = stopping_conditions self.task_pool = [] assert ( diff --git a/swarms/structs/sermon_swarm.py b/swarms/structs/sermon_swarm.py index 147ba36e..895670ad 100644 --- a/swarms/structs/sermon_swarm.py +++ b/swarms/structs/sermon_swarm.py @@ -2,6 +2,7 @@ from typing import Union, Sequence, List, Callable from swarms.structs.agent import Agent from swarms.structs.base_swarm import AbstractSwarm + class SermonSwarm(AbstractSwarm): """ Represents a swarm of agents that communicate through sermons. diff --git a/swarms/structs/task_queue_base.py b/swarms/structs/task_queue_base.py index 968023b0..b95421a5 100644 --- a/swarms/structs/task_queue_base.py +++ b/swarms/structs/task_queue_base.py @@ -1,5 +1,5 @@ import threading -from abc import ABC, abstractmethod +from abc import ABC from swarms.structs.agent import Agent from swarms.structs.task import Task @@ -36,8 +36,8 @@ class TaskQueueBase(ABC): self.lock = threading.Lock() @synchronized_queue - @abstractmethod - def add_task(self, task: Task) -> bool: + # @abstractmethod + def add(self, task: Task) -> bool: """Adds a task to the queue. Args: @@ -46,11 +46,11 @@ class TaskQueueBase(ABC): Returns: bool: True if the task was successfully added, False otherwise. """ - raise NotImplementedError + ... @synchronized_queue - @abstractmethod - def get_task(self, agent: Agent) -> Task: + # @abstractmethod + def get(self, agent: Agent) -> Task: """Gets the next task from the queue. Args: @@ -59,24 +59,24 @@ class TaskQueueBase(ABC): Returns: Task: The next task from the queue. """ - raise NotImplementedError + ... @synchronized_queue - @abstractmethod + # @abstractmethod def complete_task(self, task_id: str): """Sets the task as completed. Args: task_id (str): The ID of the task to be marked as completed. """ - raise NotImplementedError + ... @synchronized_queue - @abstractmethod - def reset_task(self, task_id: str): + # @abstractmethod + def reset(self, task_id: str): """Resets the task if the agent failed to complete it. Args: task_id (str): The ID of the task to be reset. """ - raise NotImplementedError + ...