[MIGRATION][Migrate SwarmMatcher, JsonFormer, and other packages and modules to a new package called swarms-utils to remove torch and transformers depedency] [next need to do numpy]
Kye Gomez
2 weeks ago
parent
f14b282a27
commit
e0da12be90
@ -1,256 +0,0 @@
|
||||
from typing import List, Optional, Dict, Any, Callable
|
||||
from loguru import logger
|
||||
from swarms.agents.exceptions import (
|
||||
ToolExecutionError,
|
||||
ToolValidationError,
|
||||
ToolNotFoundError,
|
||||
ToolParameterError,
|
||||
)
|
||||
|
||||
|
||||
class ToolAgent:
|
||||
"""
|
||||
A wrapper class for vLLM that provides a similar interface to LiteLLM.
|
||||
This class handles model initialization and inference using vLLM.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str = "meta-llama/Llama-2-7b-chat-hf",
|
||||
system_prompt: Optional[str] = None,
|
||||
stream: bool = False,
|
||||
temperature: float = 0.5,
|
||||
max_tokens: int = 4000,
|
||||
max_completion_tokens: int = 4000,
|
||||
tools_list_dictionary: Optional[List[Dict[str, Any]]] = None,
|
||||
tool_choice: str = "auto",
|
||||
parallel_tool_calls: bool = False,
|
||||
retry_attempts: int = 3,
|
||||
retry_interval: float = 1.0,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Initialize the vLLM wrapper with the given parameters.
|
||||
Args:
|
||||
model_name (str): The name of the model to use. Defaults to "meta-llama/Llama-2-7b-chat-hf".
|
||||
system_prompt (str, optional): The system prompt to use. Defaults to None.
|
||||
stream (bool): Whether to stream the output. Defaults to False.
|
||||
temperature (float): The temperature for sampling. Defaults to 0.5.
|
||||
max_tokens (int): The maximum number of tokens to generate. Defaults to 4000.
|
||||
max_completion_tokens (int): The maximum number of completion tokens. Defaults to 4000.
|
||||
tools_list_dictionary (List[Dict[str, Any]], optional): List of available tools. Defaults to None.
|
||||
tool_choice (str): How to choose tools. Defaults to "auto".
|
||||
parallel_tool_calls (bool): Whether to allow parallel tool calls. Defaults to False.
|
||||
retry_attempts (int): Number of retry attempts for failed operations. Defaults to 3.
|
||||
retry_interval (float): Time to wait between retries in seconds. Defaults to 1.0.
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self.system_prompt = system_prompt
|
||||
self.stream = stream
|
||||
self.temperature = temperature
|
||||
self.max_tokens = max_tokens
|
||||
self.max_completion_tokens = max_completion_tokens
|
||||
self.tools_list_dictionary = tools_list_dictionary
|
||||
self.tool_choice = tool_choice
|
||||
self.parallel_tool_calls = parallel_tool_calls
|
||||
self.retry_attempts = retry_attempts
|
||||
self.retry_interval = retry_interval
|
||||
|
||||
# Initialize vLLM
|
||||
try:
|
||||
self.llm = LLM(model=model_name, **kwargs)
|
||||
self.sampling_params = SamplingParams(
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
except Exception as e:
|
||||
raise ToolExecutionError(
|
||||
"model_initialization",
|
||||
e,
|
||||
{"model_name": model_name, "kwargs": kwargs},
|
||||
)
|
||||
|
||||
def _validate_tool(
|
||||
self, tool_name: str, parameters: Dict[str, Any]
|
||||
) -> None:
|
||||
"""
|
||||
Validate tool parameters before execution.
|
||||
Args:
|
||||
tool_name (str): Name of the tool to validate
|
||||
parameters (Dict[str, Any]): Parameters to validate
|
||||
Raises:
|
||||
ToolValidationError: If validation fails
|
||||
"""
|
||||
if not self.tools_list_dictionary:
|
||||
raise ToolValidationError(
|
||||
tool_name,
|
||||
"parameters",
|
||||
"No tools available for validation",
|
||||
)
|
||||
|
||||
tool_spec = next(
|
||||
(
|
||||
tool
|
||||
for tool in self.tools_list_dictionary
|
||||
if tool["name"] == tool_name
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
if not tool_spec:
|
||||
raise ToolNotFoundError(tool_name)
|
||||
|
||||
required_params = {
|
||||
param["name"]
|
||||
for param in tool_spec["parameters"]
|
||||
if param.get("required", True)
|
||||
}
|
||||
|
||||
missing_params = required_params - set(parameters.keys())
|
||||
if missing_params:
|
||||
raise ToolParameterError(
|
||||
tool_name,
|
||||
f"Missing required parameters: {', '.join(missing_params)}",
|
||||
)
|
||||
|
||||
def _execute_with_retry(
|
||||
self, func: Callable, *args, **kwargs
|
||||
) -> Any:
|
||||
"""
|
||||
Execute a function with retry logic.
|
||||
Args:
|
||||
func (Callable): Function to execute
|
||||
*args: Positional arguments for the function
|
||||
**kwargs: Keyword arguments for the function
|
||||
Returns:
|
||||
Any: Result of the function execution
|
||||
Raises:
|
||||
ToolExecutionError: If all retry attempts fail
|
||||
"""
|
||||
last_error = None
|
||||
for attempt in range(self.retry_attempts):
|
||||
try:
|
||||
return func(*args, **kwargs)
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
logger.warning(
|
||||
f"Attempt {attempt + 1}/{self.retry_attempts} failed: {str(e)}"
|
||||
)
|
||||
if attempt < self.retry_attempts - 1:
|
||||
time.sleep(self.retry_interval)
|
||||
|
||||
raise ToolExecutionError(
|
||||
func.__name__,
|
||||
last_error,
|
||||
{"attempts": self.retry_attempts},
|
||||
)
|
||||
|
||||
def run(self, task: str, *args, **kwargs) -> str:
|
||||
"""
|
||||
Run the tool agent for the specified task.
|
||||
Args:
|
||||
task (str): The task to be performed by the tool agent.
|
||||
*args: Variable length argument list.
|
||||
**kwargs: Arbitrary keyword arguments.
|
||||
Returns:
|
||||
The output of the tool agent.
|
||||
Raises:
|
||||
ToolExecutionError: If an error occurs during execution.
|
||||
"""
|
||||
try:
|
||||
if not self.llm:
|
||||
raise ToolExecutionError(
|
||||
"run",
|
||||
Exception("LLM not initialized"),
|
||||
{"task": task},
|
||||
)
|
||||
|
||||
logger.info(f"Running task: {task}")
|
||||
|
||||
# Prepare the prompt
|
||||
prompt = self._prepare_prompt(task)
|
||||
|
||||
# Execute with retry logic
|
||||
outputs = self._execute_with_retry(
|
||||
self.llm.generate, prompt, self.sampling_params
|
||||
)
|
||||
|
||||
response = outputs[0].outputs[0].text.strip()
|
||||
return response
|
||||
|
||||
except Exception as error:
|
||||
logger.error(f"Error running task: {error}")
|
||||
raise ToolExecutionError(
|
||||
"run",
|
||||
error,
|
||||
{"task": task, "args": args, "kwargs": kwargs},
|
||||
)
|
||||
|
||||
def _prepare_prompt(self, task: str) -> str:
|
||||
"""
|
||||
Prepare the prompt for the given task.
|
||||
Args:
|
||||
task (str): The task to prepare the prompt for.
|
||||
Returns:
|
||||
str: The prepared prompt.
|
||||
"""
|
||||
if self.system_prompt:
|
||||
return f"{self.system_prompt}\n\nUser: {task}\nAssistant:"
|
||||
return f"User: {task}\nAssistant:"
|
||||
|
||||
def __call__(self, task: str, *args, **kwargs) -> str:
|
||||
"""
|
||||
Call the model for the given task.
|
||||
Args:
|
||||
task (str): The task to run the model for.
|
||||
*args: Additional positional arguments.
|
||||
**kwargs: Additional keyword arguments.
|
||||
Returns:
|
||||
str: The model's response.
|
||||
"""
|
||||
return self.run(task, *args, **kwargs)
|
||||
|
||||
def batched_run(
|
||||
self, tasks: List[str], batch_size: int = 10
|
||||
) -> List[str]:
|
||||
"""
|
||||
Run the model for multiple tasks in batches.
|
||||
Args:
|
||||
tasks (List[str]): List of tasks to run.
|
||||
batch_size (int): Size of each batch. Defaults to 10.
|
||||
Returns:
|
||||
List[str]: List of model responses.
|
||||
Raises:
|
||||
ToolExecutionError: If an error occurs during batch execution.
|
||||
"""
|
||||
logger.info(
|
||||
f"Running tasks in batches of size {batch_size}. Total tasks: {len(tasks)}"
|
||||
)
|
||||
results = []
|
||||
|
||||
try:
|
||||
for i in range(0, len(tasks), batch_size):
|
||||
batch = tasks[i : i + batch_size]
|
||||
for task in batch:
|
||||
logger.info(f"Running task: {task}")
|
||||
try:
|
||||
result = self.run(task)
|
||||
results.append(result)
|
||||
except ToolExecutionError as e:
|
||||
logger.error(
|
||||
f"Failed to execute task '{task}': {e}"
|
||||
)
|
||||
results.append(f"Error: {str(e)}")
|
||||
continue
|
||||
|
||||
logger.info("Completed all tasks.")
|
||||
return results
|
||||
|
||||
except Exception as error:
|
||||
logger.error(f"Error in batch execution: {error}")
|
||||
raise ToolExecutionError(
|
||||
"batched_run",
|
||||
error,
|
||||
{"tasks": tasks, "batch_size": batch_size},
|
||||
)
|
||||
@ -1,599 +0,0 @@
|
||||
import json
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
from pydantic import BaseModel, Field
|
||||
from tenacity import retry, stop_after_attempt, wait_exponential
|
||||
|
||||
from swarms.utils.loguru_logger import initialize_logger
|
||||
|
||||
logger = initialize_logger(log_folder="swarm_matcher")
|
||||
|
||||
|
||||
class SwarmType(BaseModel):
|
||||
name: str
|
||||
description: str
|
||||
embedding: Optional[List[float]] = Field(
|
||||
default=None, exclude=True
|
||||
)
|
||||
|
||||
|
||||
class SwarmMatcherConfig(BaseModel):
|
||||
model_name: str = "sentence-transformers/all-MiniLM-L6-v2"
|
||||
embedding_dim: int = (
|
||||
512 # Dimension of the sentence-transformers model
|
||||
)
|
||||
|
||||
|
||||
class SwarmMatcher:
|
||||
"""
|
||||
A class for matching tasks to swarm types based on their descriptions.
|
||||
It utilizes a transformer model to generate embeddings for task and swarm type descriptions,
|
||||
and then calculates the dot product to find the best match.
|
||||
"""
|
||||
|
||||
def __init__(self, config: SwarmMatcherConfig):
|
||||
"""
|
||||
Initializes the SwarmMatcher with a configuration.
|
||||
|
||||
Args:
|
||||
config (SwarmMatcherConfig): The configuration for the SwarmMatcher.
|
||||
"""
|
||||
logger.add("swarm_matcher_debug.log", level="DEBUG")
|
||||
logger.debug("Initializing SwarmMatcher")
|
||||
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"torch package not found. Pip install torch."
|
||||
)
|
||||
|
||||
try:
|
||||
import transformers
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"transformers package not found. Pip install transformers."
|
||||
)
|
||||
|
||||
self.torch = torch
|
||||
try:
|
||||
self.config = config
|
||||
self.tokenizer = (
|
||||
transformers.AutoTokenizer.from_pretrained(
|
||||
config.model_name
|
||||
)
|
||||
)
|
||||
self.model = transformers.AutoModel.from_pretrained(
|
||||
config.model_name
|
||||
)
|
||||
self.swarm_types: List[SwarmType] = []
|
||||
logger.debug("SwarmMatcher initialized successfully")
|
||||
except Exception as e:
|
||||
logger.error(f"Error initializing SwarmMatcher: {str(e)}")
|
||||
raise
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
wait=wait_exponential(multiplier=1, min=4, max=10),
|
||||
)
|
||||
def get_embedding(self, text: str) -> np.ndarray:
|
||||
"""
|
||||
Generates an embedding for a given text using the configured model.
|
||||
|
||||
Args:
|
||||
text (str): The text for which to generate an embedding.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The embedding vector for the text.
|
||||
"""
|
||||
logger.debug(f"Getting embedding for text: {text[:50]}...")
|
||||
try:
|
||||
inputs = self.tokenizer(
|
||||
text,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=512,
|
||||
)
|
||||
with self.torch.no_grad():
|
||||
outputs = self.model(**inputs)
|
||||
embedding = (
|
||||
outputs.last_hidden_state.mean(dim=1)
|
||||
.squeeze()
|
||||
.numpy()
|
||||
)
|
||||
logger.debug("Embedding generated successfully")
|
||||
return embedding
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating embedding: {str(e)}")
|
||||
raise
|
||||
|
||||
def add_swarm_type(self, swarm_type: SwarmType):
|
||||
"""
|
||||
Adds a swarm type to the list of swarm types, generating an embedding for its description.
|
||||
|
||||
Args:
|
||||
swarm_type (SwarmType): The swarm type to add.
|
||||
"""
|
||||
logger.debug(f"Adding swarm type: {swarm_type.name}")
|
||||
try:
|
||||
embedding = self.get_embedding(swarm_type.description)
|
||||
swarm_type.embedding = embedding.tolist()
|
||||
self.swarm_types.append(swarm_type)
|
||||
logger.info(f"Added swarm type: {swarm_type.name}")
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error adding swarm type {swarm_type.name}: {str(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
def find_best_match(self, task: str) -> Tuple[str, float]:
|
||||
"""
|
||||
Finds the best match for a given task among the registered swarm types.
|
||||
|
||||
Args:
|
||||
task (str): The task for which to find the best match.
|
||||
|
||||
Returns:
|
||||
Tuple[str, float]: A tuple containing the name of the best matching swarm type and the score.
|
||||
"""
|
||||
logger.debug(f"Finding best match for task: {task[:50]}...")
|
||||
try:
|
||||
task_embedding = self.get_embedding(task)
|
||||
best_match = None
|
||||
best_score = -float("inf")
|
||||
for swarm_type in self.swarm_types:
|
||||
score = np.dot(
|
||||
task_embedding, np.array(swarm_type.embedding)
|
||||
)
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_match = swarm_type
|
||||
logger.info(
|
||||
f"Best match for task: {best_match.name} (score: {best_score})"
|
||||
)
|
||||
return best_match.name, float(best_score)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error finding best match for task: {str(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
def auto_select_swarm(self, task: str) -> str:
|
||||
"""
|
||||
Automatically selects the best swarm type for a given task based on their descriptions.
|
||||
|
||||
Args:
|
||||
task (str): The task for which to select a swarm type.
|
||||
|
||||
Returns:
|
||||
str: The name of the selected swarm type.
|
||||
"""
|
||||
logger.debug(f"Auto-selecting swarm for task: {task[:50]}...")
|
||||
best_match, score = self.find_best_match(task)
|
||||
logger.info(f"Task: {task}")
|
||||
logger.info(f"Selected Swarm Type: {best_match}")
|
||||
logger.info(f"Confidence Score: {score:.2f}")
|
||||
return best_match
|
||||
|
||||
def run_multiple(self, tasks: List[str], *args, **kwargs) -> str:
|
||||
swarms = []
|
||||
|
||||
for task in tasks:
|
||||
output = self.auto_select_swarm(task)
|
||||
|
||||
# Append
|
||||
swarms.append(output)
|
||||
|
||||
return swarms
|
||||
|
||||
def save_swarm_types(self, filename: str):
|
||||
"""
|
||||
Saves the registered swarm types to a JSON file.
|
||||
|
||||
Args:
|
||||
filename (str): The name of the file to which to save the swarm types.
|
||||
"""
|
||||
try:
|
||||
with open(filename, "w") as f:
|
||||
json.dump([st.dict() for st in self.swarm_types], f)
|
||||
logger.info(f"Saved swarm types to {filename}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error saving swarm types: {str(e)}")
|
||||
raise
|
||||
|
||||
def load_swarm_types(self, filename: str):
|
||||
"""
|
||||
Loads swarm types from a JSON file.
|
||||
|
||||
Args:
|
||||
filename (str): The name of the file from which to load the swarm types.
|
||||
"""
|
||||
try:
|
||||
with open(filename, "r") as f:
|
||||
swarm_types_data = json.load(f)
|
||||
self.swarm_types = [
|
||||
SwarmType(**st) for st in swarm_types_data
|
||||
]
|
||||
logger.info(f"Loaded swarm types from {filename}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading swarm types: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
def initialize_swarm_types(matcher: SwarmMatcher):
|
||||
logger.debug("Initializing swarm types")
|
||||
swarm_types = [
|
||||
SwarmType(
|
||||
name="AgentRearrange",
|
||||
description="Optimize agent order and rearrange flow for multi-step tasks, ensuring efficient task allocation and minimizing bottlenecks. Keywords: orchestration, coordination, pipeline optimization, task scheduling, resource allocation, workflow management, agent organization, process optimization",
|
||||
),
|
||||
SwarmType(
|
||||
name="MixtureOfAgents",
|
||||
description="Combine diverse expert agents for comprehensive analysis, fostering a collaborative approach to problem-solving and leveraging individual strengths. Keywords: multi-agent system, expert collaboration, distributed intelligence, collective problem solving, agent specialization, team coordination, hybrid approaches, knowledge synthesis",
|
||||
),
|
||||
SwarmType(
|
||||
name="SpreadSheetSwarm",
|
||||
description="Collaborative data processing and analysis in a spreadsheet-like environment, facilitating real-time data sharing and visualization. Keywords: data analysis, tabular processing, collaborative editing, data transformation, spreadsheet operations, data visualization, real-time collaboration, structured data",
|
||||
),
|
||||
SwarmType(
|
||||
name="SequentialWorkflow",
|
||||
description="Execute tasks in a step-by-step, sequential process workflow, ensuring a logical and methodical approach to task execution. Keywords: linear processing, waterfall methodology, step-by-step execution, ordered tasks, sequential operations, process flow, systematic approach, staged execution",
|
||||
),
|
||||
SwarmType(
|
||||
name="ConcurrentWorkflow",
|
||||
description="Process multiple tasks or data sources concurrently in parallel, maximizing productivity and reducing processing time. Keywords: parallel processing, multi-threading, asynchronous execution, distributed computing, concurrent operations, simultaneous tasks, parallel workflows, scalable processing",
|
||||
),
|
||||
SwarmType(
|
||||
name="HierarchicalSwarm",
|
||||
description="Organize agents in a hierarchical structure with clear reporting lines and delegation of responsibilities. Keywords: management hierarchy, organizational structure, delegation, supervision, chain of command, tiered organization, structured coordination",
|
||||
),
|
||||
# SwarmType(
|
||||
# name="AdaptiveSwarm",
|
||||
# description="Dynamically adjust agent behavior and swarm configuration based on task requirements and performance feedback. Keywords: dynamic adaptation, self-optimization, feedback loops, learning systems, flexible configuration, responsive behavior, adaptive algorithms",
|
||||
# ),
|
||||
# SwarmType(
|
||||
# name="ConsensusSwarm",
|
||||
# description="Achieve group decisions through consensus mechanisms and voting protocols among multiple agents. Keywords: group decision making, voting systems, collective intelligence, agreement protocols, democratic processes, collaborative decisions",
|
||||
# ),
|
||||
]
|
||||
|
||||
for swarm_type in swarm_types:
|
||||
matcher.add_swarm_type(swarm_type)
|
||||
logger.debug("Swarm types initialized")
|
||||
|
||||
|
||||
def swarm_matcher(task: str, *args, **kwargs):
|
||||
"""
|
||||
Runs the SwarmMatcher example with predefined tasks and swarm types.
|
||||
"""
|
||||
config = SwarmMatcherConfig()
|
||||
matcher = SwarmMatcher(config)
|
||||
initialize_swarm_types(matcher)
|
||||
|
||||
# matcher.save_swarm_types(f"swarm_logs/{uuid4().hex}.json")
|
||||
|
||||
swarm_type = matcher.auto_select_swarm(task)
|
||||
|
||||
logger.info(f"{swarm_type}")
|
||||
|
||||
return swarm_type
|
||||
|
||||
|
||||
# from typing import List, Tuple, Dict
|
||||
# from pydantic import BaseModel, Field
|
||||
# from loguru import logger
|
||||
# from uuid import uuid4
|
||||
# import chromadb
|
||||
# import json
|
||||
# from tenacity import retry, stop_after_attempt, wait_exponential
|
||||
|
||||
|
||||
# class SwarmType(BaseModel):
|
||||
# """A swarm type with its name, description and optional metadata"""
|
||||
|
||||
# id: str = Field(default_factory=lambda: str(uuid4()))
|
||||
# name: str
|
||||
# description: str
|
||||
# metadata: Dict = Field(default_factory=dict)
|
||||
|
||||
|
||||
# class SwarmMatcherConfig(BaseModel):
|
||||
# """Configuration for the SwarmMatcher"""
|
||||
|
||||
# collection_name: str = "swarm_types"
|
||||
# distance_metric: str = "cosine" # or "l2" or "ip"
|
||||
# embedding_function: str = (
|
||||
# "sentence-transformers/all-mpnet-base-v2" # Better model than MiniLM
|
||||
# )
|
||||
# persist_directory: str = "./chroma_db"
|
||||
|
||||
|
||||
# class SwarmMatcher:
|
||||
# """
|
||||
# An improved swarm matcher that uses ChromaDB for better vector similarity search.
|
||||
# Features:
|
||||
# - Persistent storage of embeddings
|
||||
# - Better vector similarity search with multiple distance metrics
|
||||
# - Improved embedding model
|
||||
# - Metadata filtering capabilities
|
||||
# - Batch operations support
|
||||
# """
|
||||
|
||||
# def __init__(self, config: SwarmMatcherConfig):
|
||||
# """Initialize the improved swarm matcher"""
|
||||
# logger.add("swarm_matcher.log", rotation="100 MB")
|
||||
# self.config = config
|
||||
|
||||
# # Initialize ChromaDB client with persistence
|
||||
# self.chroma_client = chromadb.Client()
|
||||
|
||||
# # Get or create collection
|
||||
# try:
|
||||
# self.collection = self.chroma_client.get_collection(
|
||||
# name=config.collection_name,
|
||||
# )
|
||||
# except ValueError:
|
||||
# self.collection = self.chroma_client.create_collection(
|
||||
# name=config.collection_name,
|
||||
# metadata={"hnsw:space": config.distance_metric},
|
||||
# )
|
||||
|
||||
# logger.info(
|
||||
# f"Initialized SwarmMatcher with collection '{config.collection_name}'"
|
||||
# )
|
||||
|
||||
# def add_swarm_type(self, swarm_type: SwarmType) -> None:
|
||||
# """Add a single swarm type to the collection"""
|
||||
# try:
|
||||
# self.collection.add(
|
||||
# ids=[swarm_type.id],
|
||||
# documents=[swarm_type.description],
|
||||
# metadatas=[
|
||||
# {"name": swarm_type.name, **swarm_type.metadata}
|
||||
# ],
|
||||
# )
|
||||
# logger.info(f"Added swarm type: {swarm_type.name}")
|
||||
# except Exception as e:
|
||||
# logger.error(
|
||||
# f"Error adding swarm type {swarm_type.name}: {str(e)}"
|
||||
# )
|
||||
# raise
|
||||
|
||||
# def add_swarm_types(self, swarm_types: List[SwarmType]) -> None:
|
||||
# """Add multiple swarm types in batch"""
|
||||
# try:
|
||||
# self.collection.add(
|
||||
# ids=[st.id for st in swarm_types],
|
||||
# documents=[st.description for st in swarm_types],
|
||||
# metadatas=[
|
||||
# {"name": st.name, **st.metadata}
|
||||
# for st in swarm_types
|
||||
# ],
|
||||
# )
|
||||
# logger.info(f"Added {len(swarm_types)} swarm types")
|
||||
# except Exception as e:
|
||||
# logger.error(
|
||||
# f"Error adding swarm types in batch: {str(e)}"
|
||||
# )
|
||||
# raise
|
||||
|
||||
# @retry(
|
||||
# stop=stop_after_attempt(3),
|
||||
# wait=wait_exponential(multiplier=1, min=4, max=10),
|
||||
# )
|
||||
# def find_best_matches(
|
||||
# self,
|
||||
# task: str,
|
||||
# n_results: int = 3,
|
||||
# score_threshold: float = 0.7,
|
||||
# ) -> List[Tuple[str, float]]:
|
||||
# """
|
||||
# Find the best matching swarm types for a given task
|
||||
# Returns multiple matches with their scores
|
||||
# """
|
||||
# try:
|
||||
# results = self.collection.query(
|
||||
# query_texts=[task],
|
||||
# n_results=n_results,
|
||||
# include=["metadatas", "distances"],
|
||||
# )
|
||||
|
||||
# matches = []
|
||||
# for metadata, distance in zip(
|
||||
# results["metadatas"][0], results["distances"][0]
|
||||
# ):
|
||||
# # Convert distance to similarity score (1 - normalized_distance)
|
||||
# score = 1 - (
|
||||
# distance / 2
|
||||
# ) # Normalize cosine distance to [0,1]
|
||||
# if score >= score_threshold:
|
||||
# matches.append((metadata["name"], score))
|
||||
|
||||
# logger.info(f"Found {len(matches)} matches for task")
|
||||
# return matches
|
||||
|
||||
# except Exception as e:
|
||||
# logger.error(f"Error finding matches for task: {str(e)}")
|
||||
# raise
|
||||
|
||||
# def auto_select_swarm(self, task: str) -> str:
|
||||
# """
|
||||
# Automatically select the best swarm type for a task
|
||||
# Returns only the top match
|
||||
# """
|
||||
# matches = self.find_best_matches(task, n_results=1)
|
||||
# if not matches:
|
||||
# logger.warning("No suitable matches found for task")
|
||||
# return "SequentialWorkflow" # Default fallback
|
||||
|
||||
# best_match, score = matches[0]
|
||||
# logger.info(
|
||||
# f"Selected swarm type '{best_match}' with confidence {score:.3f}"
|
||||
# )
|
||||
# return best_match
|
||||
|
||||
# def run_multiple(self, tasks: List[str]) -> List[str]:
|
||||
# """Process multiple tasks in batch"""
|
||||
# return [self.auto_select_swarm(task) for task in tasks]
|
||||
|
||||
# def save_swarm_types(self, filename: str) -> None:
|
||||
# """Export swarm types to JSON"""
|
||||
# try:
|
||||
# all_data = self.collection.get(
|
||||
# include=["metadatas", "documents"]
|
||||
# )
|
||||
# swarm_types = [
|
||||
# SwarmType(
|
||||
# id=id_,
|
||||
# name=metadata["name"],
|
||||
# description=document,
|
||||
# metadata={
|
||||
# k: v
|
||||
# for k, v in metadata.items()
|
||||
# if k != "name"
|
||||
# },
|
||||
# )
|
||||
# for id_, metadata, document in zip(
|
||||
# all_data["ids"],
|
||||
# all_data["metadatas"],
|
||||
# all_data["documents"],
|
||||
# )
|
||||
# ]
|
||||
|
||||
# with open(filename, "w") as f:
|
||||
# json.dump(
|
||||
# [st.dict() for st in swarm_types], f, indent=2
|
||||
# )
|
||||
# logger.info(f"Saved swarm types to {filename}")
|
||||
# except Exception as e:
|
||||
# logger.error(f"Error saving swarm types: {str(e)}")
|
||||
# raise
|
||||
|
||||
# def load_swarm_types(self, filename: str) -> None:
|
||||
# """Import swarm types from JSON"""
|
||||
# try:
|
||||
# with open(filename, "r") as f:
|
||||
# swarm_types_data = json.load(f)
|
||||
# swarm_types = [SwarmType(**st) for st in swarm_types_data]
|
||||
# self.add_swarm_types(swarm_types)
|
||||
# logger.info(f"Loaded swarm types from {filename}")
|
||||
# except Exception as e:
|
||||
# logger.error(f"Error loading swarm types: {str(e)}")
|
||||
# raise
|
||||
|
||||
|
||||
# def initialize_default_swarm_types(matcher: SwarmMatcher) -> None:
|
||||
# """Initialize the matcher with default swarm types"""
|
||||
# swarm_types = [
|
||||
# SwarmType(
|
||||
# name="AgentRearrange",
|
||||
# description="""
|
||||
# Optimize agent order and rearrange flow for multi-step tasks, ensuring efficient task allocation
|
||||
# and minimizing bottlenecks. Specialized in orchestration, coordination, pipeline optimization,
|
||||
# task scheduling, resource allocation, workflow management, agent organization, and process optimization.
|
||||
# Best for tasks requiring complex agent interactions and workflow optimization.
|
||||
# """,
|
||||
# metadata={
|
||||
# "category": "optimization",
|
||||
# "complexity": "high",
|
||||
# },
|
||||
# ),
|
||||
# SwarmType(
|
||||
# name="MixtureOfAgents",
|
||||
# description="""
|
||||
# Combine diverse expert agents for comprehensive analysis, fostering a collaborative approach
|
||||
# to problem-solving and leveraging individual strengths. Focuses on multi-agent systems,
|
||||
# expert collaboration, distributed intelligence, collective problem solving, agent specialization,
|
||||
# team coordination, hybrid approaches, and knowledge synthesis. Ideal for complex problems
|
||||
# requiring multiple areas of expertise.
|
||||
# """,
|
||||
# metadata={
|
||||
# "category": "collaboration",
|
||||
# "complexity": "high",
|
||||
# },
|
||||
# ),
|
||||
# SwarmType(
|
||||
# name="SpreadSheetSwarm",
|
||||
# description="""
|
||||
# Collaborative data processing and analysis in a spreadsheet-like environment, facilitating
|
||||
# real-time data sharing and visualization. Specializes in data analysis, tabular processing,
|
||||
# collaborative editing, data transformation, spreadsheet operations, data visualization,
|
||||
# real-time collaboration, and structured data handling. Perfect for data-intensive tasks
|
||||
# requiring structured analysis.
|
||||
# """,
|
||||
# metadata={
|
||||
# "category": "data_processing",
|
||||
# "complexity": "medium",
|
||||
# },
|
||||
# ),
|
||||
# SwarmType(
|
||||
# name="SequentialWorkflow",
|
||||
# description="""
|
||||
# Execute tasks in a step-by-step, sequential process workflow, ensuring a logical and methodical
|
||||
# approach to task execution. Focuses on linear processing, waterfall methodology, step-by-step
|
||||
# execution, ordered tasks, sequential operations, process flow, systematic approach, and staged
|
||||
# execution. Best for tasks requiring strict order and dependencies.
|
||||
# """,
|
||||
# metadata={"category": "workflow", "complexity": "low"},
|
||||
# ),
|
||||
# SwarmType(
|
||||
# name="ConcurrentWorkflow",
|
||||
# description="""
|
||||
# Process multiple tasks or data sources concurrently in parallel, maximizing productivity
|
||||
# and reducing processing time. Specializes in parallel processing, multi-threading,
|
||||
# asynchronous execution, distributed computing, concurrent operations, simultaneous tasks,
|
||||
# parallel workflows, and scalable processing. Ideal for independent tasks that can be
|
||||
# processed simultaneously.
|
||||
# """,
|
||||
# metadata={"category": "workflow", "complexity": "medium"},
|
||||
# ),
|
||||
# ]
|
||||
|
||||
# matcher.add_swarm_types(swarm_types)
|
||||
# logger.info("Initialized default swarm types")
|
||||
|
||||
|
||||
# def create_swarm_matcher(
|
||||
# persist_dir: str = "./chroma_db",
|
||||
# collection_name: str = "swarm_types",
|
||||
# ) -> SwarmMatcher:
|
||||
# """Convenience function to create and initialize a swarm matcher"""
|
||||
# config = SwarmMatcherConfig(
|
||||
# persist_directory=persist_dir, collection_name=collection_name
|
||||
# )
|
||||
# matcher = SwarmMatcher(config)
|
||||
# initialize_default_swarm_types(matcher)
|
||||
# return matcher
|
||||
|
||||
|
||||
# # Example usage
|
||||
# def swarm_matcher(task: str) -> str:
|
||||
# # Create and initialize matcher
|
||||
# matcher = create_swarm_matcher()
|
||||
|
||||
# swarm_type = matcher.auto_select_swarm(task)
|
||||
# print(f"Task: {task}\nSelected Swarm: {swarm_type}\n")
|
||||
|
||||
# return swarm_type
|
||||
|
||||
|
||||
# # # Example usage
|
||||
# # if __name__ == "__main__":
|
||||
# # # Create and initialize matcher
|
||||
# # matcher = create_swarm_matcher()
|
||||
|
||||
# # # Example tasks
|
||||
# # tasks = [
|
||||
# # "Analyze this spreadsheet of sales data and create visualizations",
|
||||
# # "Coordinate multiple AI agents to solve a complex problem",
|
||||
# # "Process these tasks one after another in a specific order",
|
||||
# # "Write multiple blog posts about the latest advancements in swarm intelligence all at once",
|
||||
# # "Write a blog post about the latest advancements in swarm intelligence",
|
||||
# # ]
|
||||
|
||||
# # # Process tasks
|
||||
# # for task in tasks:
|
||||
# # swarm_type = matcher.auto_select_swarm(task)
|
||||
# # print(f"Task: {task}\nSelected Swarm: {swarm_type}\n")
|
||||
@ -1,421 +0,0 @@
|
||||
import json
|
||||
from typing import Any, Dict, List, Union
|
||||
|
||||
from pydantic import BaseModel
|
||||
from swarms.tools.logits_processor import (
|
||||
NumberStoppingCriteria,
|
||||
OutputNumbersTokens,
|
||||
StringStoppingCriteria,
|
||||
)
|
||||
from swarms.utils.auto_download_check_packages import (
|
||||
auto_check_and_download_package,
|
||||
)
|
||||
|
||||
try:
|
||||
import transformers
|
||||
except ImportError:
|
||||
auto_check_and_download_package(
|
||||
"transformers", package_manager="pip"
|
||||
)
|
||||
import transformers
|
||||
|
||||
|
||||
GENERATION_MARKER = "|GENERATION|"
|
||||
|
||||
|
||||
class Jsonformer:
|
||||
"""
|
||||
Initializes the FormatTools class.
|
||||
|
||||
Args:
|
||||
model (PreTrainedModel): The pre-trained model.
|
||||
tokenizer (PreTrainedTokenizer): The tokenizer for the model.
|
||||
json_schema (Dict[str, Any]): The JSON schema.
|
||||
prompt (str): The prompt for generation.
|
||||
|
||||
Keyword Args:
|
||||
debug (bool, optional): Whether to enable debug mode. Defaults to False.
|
||||
max_array_length (int, optional): The maximum length of an array. Defaults to 10.
|
||||
max_number_tokens (int, optional): The maximum number of tokens for numbers. Defaults to 6.
|
||||
temperature (float, optional): The temperature for generation. Defaults to 1.0.
|
||||
max_string_token_length (int, optional): The maximum length of a string token. Defaults to 10.
|
||||
"""
|
||||
|
||||
value: Dict[str, Any] = {}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: transformers.PreTrainedModel = None, # type: ignore
|
||||
tokenizer: transformers.PreTrainedTokenizer = None, # type: ignore
|
||||
json_schema: Union[Dict[str, Any], BaseModel] = None,
|
||||
schemas: List[Union[Dict[str, Any], BaseModel]] = [],
|
||||
prompt: str = None,
|
||||
*,
|
||||
debug: bool = False,
|
||||
max_array_length: int = 10,
|
||||
max_number_tokens: int = 6,
|
||||
temperature: float = 1.0,
|
||||
max_string_token_length: int = 10,
|
||||
llm: Any = None,
|
||||
):
|
||||
self.model = model
|
||||
self.tokenizer = tokenizer
|
||||
self.json_schema = json_schema
|
||||
self.prompt = prompt
|
||||
self.llm = llm
|
||||
self.schemas = schemas
|
||||
|
||||
self.number_logit_processor = OutputNumbersTokens(
|
||||
self.tokenizer, self.prompt
|
||||
)
|
||||
|
||||
self.generation_marker = "|GENERATION|"
|
||||
self.debug_on = debug
|
||||
self.max_array_length = max_array_length
|
||||
|
||||
self.max_number_tokens = max_number_tokens
|
||||
self.temperature = temperature
|
||||
self.max_string_token_length = max_string_token_length
|
||||
|
||||
def generate_number(
|
||||
self, temperature: Union[float, None] = None, iterations=0
|
||||
):
|
||||
"""
|
||||
Generates a number based on the given prompt.
|
||||
|
||||
Args:
|
||||
temperature (float, optional): The temperature value for number generation. Defaults to None.
|
||||
iterations (int, optional): The number of iterations for generating a valid number. Defaults to 0.
|
||||
|
||||
Returns:
|
||||
float: The generated number.
|
||||
|
||||
Raises:
|
||||
ValueError: If a valid number cannot be generated after 3 iterations.
|
||||
"""
|
||||
if self.model:
|
||||
prompt = self.get_prompt()
|
||||
self.debug("[generate_number]", prompt, is_prompt=True)
|
||||
input_tokens = self.tokenizer.encode(
|
||||
prompt, return_tensors="pt"
|
||||
).to(self.model.device)
|
||||
|
||||
response = self.model.generate(
|
||||
input_tokens,
|
||||
max_new_tokens=self.max_number_tokens,
|
||||
num_return_sequences=1,
|
||||
logits_processor=[self.number_logit_processor],
|
||||
stopping_criteria=[
|
||||
NumberStoppingCriteria(
|
||||
self.tokenizer, len(input_tokens[0])
|
||||
)
|
||||
],
|
||||
temperature=temperature or self.temperature,
|
||||
pad_token_id=self.tokenizer.eos_token_id,
|
||||
)
|
||||
response = self.tokenizer.decode(
|
||||
response[0], skip_special_tokens=True
|
||||
)
|
||||
|
||||
response = response[len(prompt) :]
|
||||
response = response.strip().rstrip(".")
|
||||
self.debug("[generate_number]", response)
|
||||
try:
|
||||
return float(response)
|
||||
except ValueError:
|
||||
if iterations > 3:
|
||||
raise ValueError(
|
||||
"Failed to generate a valid number"
|
||||
)
|
||||
|
||||
return self.generate_number(
|
||||
temperature=self.temperature * 1.3,
|
||||
iterations=iterations + 1,
|
||||
)
|
||||
elif self.llm:
|
||||
prompt = self.get_prompt()
|
||||
self.debug("[generate_number]", prompt, is_prompt=True)
|
||||
response = self.llm(prompt)
|
||||
response = response[len(prompt) :]
|
||||
response = response.strip().rstrip(".")
|
||||
self.debug("[generate_number]", response)
|
||||
try:
|
||||
return float(response)
|
||||
except ValueError:
|
||||
if iterations > 3:
|
||||
raise ValueError(
|
||||
"Failed to generate a valid number"
|
||||
)
|
||||
|
||||
return self.generate_number(
|
||||
temperature=self.temperature * 1.3,
|
||||
iterations=iterations + 1,
|
||||
)
|
||||
|
||||
elif self.llm and self.model:
|
||||
raise ValueError("Both LLM and model cannot be None")
|
||||
|
||||
def generate_boolean(self) -> bool:
|
||||
"""
|
||||
Generates a boolean value based on the given prompt.
|
||||
|
||||
Returns:
|
||||
bool: The generated boolean value.
|
||||
"""
|
||||
if self.model:
|
||||
prompt = self.get_prompt()
|
||||
self.debug("[generate_boolean]", prompt, is_prompt=True)
|
||||
|
||||
input_tensor = self.tokenizer.encode(
|
||||
prompt, return_tensors="pt"
|
||||
)
|
||||
output = self.model.forward(
|
||||
input_tensor.to(self.model.device)
|
||||
)
|
||||
logits = output.logits[0, -1]
|
||||
|
||||
# todo: this assumes that "true" and "false" are both tokenized to a single token
|
||||
# this is probably not true for all tokenizers
|
||||
# this can be fixed by looking at only the first token of both "true" and "false"
|
||||
true_token_id = self.tokenizer.convert_tokens_to_ids(
|
||||
"true"
|
||||
)
|
||||
false_token_id = self.tokenizer.convert_tokens_to_ids(
|
||||
"false"
|
||||
)
|
||||
|
||||
result = logits[true_token_id] > logits[false_token_id]
|
||||
|
||||
self.debug("[generate_boolean]", result)
|
||||
|
||||
return result.item()
|
||||
|
||||
elif self.llm:
|
||||
prompt = self.get_prompt()
|
||||
self.debug("[generate_boolean]", prompt, is_prompt=True)
|
||||
|
||||
output = self.llm(prompt)
|
||||
|
||||
return output if output == "true" or "false" else None
|
||||
|
||||
else:
|
||||
raise ValueError("Both LLM and model cannot be None")
|
||||
|
||||
def generate_string(self) -> str:
|
||||
if self.model:
|
||||
prompt = self.get_prompt() + '"'
|
||||
self.debug("[generate_string]", prompt, is_prompt=True)
|
||||
input_tokens = self.tokenizer.encode(
|
||||
prompt, return_tensors="pt"
|
||||
).to(self.model.device)
|
||||
|
||||
response = self.model.generate(
|
||||
input_tokens,
|
||||
max_new_tokens=self.max_string_token_length,
|
||||
num_return_sequences=1,
|
||||
temperature=self.temperature,
|
||||
stopping_criteria=[
|
||||
StringStoppingCriteria(
|
||||
self.tokenizer, len(input_tokens[0])
|
||||
)
|
||||
],
|
||||
pad_token_id=self.tokenizer.eos_token_id,
|
||||
)
|
||||
|
||||
# Some models output the prompt as part of the response
|
||||
# This removes the prompt from the response if it is present
|
||||
if (
|
||||
len(response[0]) >= len(input_tokens[0])
|
||||
and (
|
||||
response[0][: len(input_tokens[0])]
|
||||
== input_tokens
|
||||
).all()
|
||||
):
|
||||
response = response[0][len(input_tokens[0]) :]
|
||||
if response.shape[0] == 1:
|
||||
response = response[0]
|
||||
|
||||
response = self.tokenizer.decode(
|
||||
response, skip_special_tokens=True
|
||||
)
|
||||
|
||||
self.debug("[generate_string]", "|" + response + "|")
|
||||
|
||||
if response.count('"') < 1:
|
||||
return response
|
||||
|
||||
return response.split('"')[0].strip()
|
||||
|
||||
elif self.llm:
|
||||
prompt = self.get_prompt() + '"'
|
||||
self.debug("[generate_string]", prompt, is_prompt=True)
|
||||
|
||||
response = self.llm(prompt)
|
||||
|
||||
# Some models output the prompt as part of the response
|
||||
# This removes the prompt from the response if it is present
|
||||
if (
|
||||
len(response[0]) >= len(input_tokens[0])
|
||||
and (
|
||||
response[0][: len(input_tokens[0])]
|
||||
== input_tokens
|
||||
).all()
|
||||
):
|
||||
response = response[0][len(input_tokens[0]) :]
|
||||
if response.shape[0] == 1:
|
||||
response = response[0]
|
||||
|
||||
self.debug("[generate_string]", "|" + response + "|")
|
||||
|
||||
if response.count('"') < 1:
|
||||
return response
|
||||
|
||||
return response.split('"')[0].strip()
|
||||
|
||||
else:
|
||||
raise ValueError("Both LLM and model cannot be None")
|
||||
|
||||
def generate_object(
|
||||
self, properties: Dict[str, Any], obj: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
for key, schema in properties.items():
|
||||
self.debug("[generate_object] generating value for", key)
|
||||
obj[key] = self.generate_value(schema, obj, key)
|
||||
return obj
|
||||
|
||||
def generate_value(
|
||||
self,
|
||||
schema: Dict[str, Any],
|
||||
obj: Union[Dict[str, Any], List[Any]],
|
||||
key: Union[str, None] = None,
|
||||
) -> Any:
|
||||
schema_type = schema["type"]
|
||||
if schema_type == "number":
|
||||
if key:
|
||||
obj[key] = self.generation_marker
|
||||
else:
|
||||
obj.append(self.generation_marker)
|
||||
return self.generate_number()
|
||||
elif schema_type == "boolean":
|
||||
if key:
|
||||
obj[key] = self.generation_marker
|
||||
else:
|
||||
obj.append(self.generation_marker)
|
||||
return self.generate_boolean()
|
||||
elif schema_type == "string":
|
||||
if key:
|
||||
obj[key] = self.generation_marker
|
||||
else:
|
||||
obj.append(self.generation_marker)
|
||||
return self.generate_string()
|
||||
elif schema_type == "array":
|
||||
new_array = []
|
||||
obj[key] = new_array
|
||||
return self.generate_array(schema["items"], new_array)
|
||||
elif schema_type == "object":
|
||||
new_obj = {}
|
||||
if key:
|
||||
obj[key] = new_obj
|
||||
else:
|
||||
obj.append(new_obj)
|
||||
return self.generate_object(schema["properties"], new_obj)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported schema type: {schema_type}"
|
||||
)
|
||||
|
||||
def generate_array(
|
||||
self, item_schema: Dict[str, Any], obj: Dict[str, Any]
|
||||
) -> list:
|
||||
if self.model:
|
||||
for _ in range(self.max_array_length):
|
||||
# forces array to have at least one element
|
||||
element = self.generate_value(item_schema, obj)
|
||||
obj[-1] = element
|
||||
|
||||
obj.append(self.generation_marker)
|
||||
input_prompt = self.get_prompt()
|
||||
obj.pop()
|
||||
input_tensor = self.tokenizer.encode(
|
||||
input_prompt, return_tensors="pt"
|
||||
)
|
||||
output = self.model.forward(
|
||||
input_tensor.to(self.model.device)
|
||||
)
|
||||
logits = output.logits[0, -1]
|
||||
|
||||
top_indices = logits.topk(30).indices
|
||||
sorted_token_ids = top_indices[
|
||||
logits[top_indices].argsort(descending=True)
|
||||
]
|
||||
|
||||
found_comma = False
|
||||
found_close_bracket = False
|
||||
|
||||
for token_id in sorted_token_ids:
|
||||
decoded_token = self.tokenizer.decode(token_id)
|
||||
if "," in decoded_token:
|
||||
found_comma = True
|
||||
break
|
||||
if "]" in decoded_token:
|
||||
found_close_bracket = True
|
||||
break
|
||||
|
||||
if found_close_bracket or not found_comma:
|
||||
break
|
||||
|
||||
return obj
|
||||
|
||||
elif self.llm:
|
||||
for _ in range(self.max_array_length):
|
||||
# forces array to have at least one element
|
||||
element = self.generate_value(item_schema, obj)
|
||||
obj[-1] = element
|
||||
|
||||
obj.append(self.generation_marker)
|
||||
input_prompt = self.get_prompt()
|
||||
obj.pop()
|
||||
output = self.llm(input_prompt)
|
||||
|
||||
found_comma = False
|
||||
found_close_bracket = False
|
||||
|
||||
for token_id in output:
|
||||
decoded_token = str(token_id)
|
||||
if "," in decoded_token:
|
||||
found_comma = True
|
||||
break
|
||||
if "]" in decoded_token:
|
||||
found_close_bracket = True
|
||||
break
|
||||
|
||||
if found_close_bracket or not found_comma:
|
||||
break
|
||||
|
||||
return obj
|
||||
|
||||
def get_prompt(self):
|
||||
template = """{prompt}\nOutput result in the following JSON schema format:\n{schema}\nResult: {progress}"""
|
||||
progress = json.dumps(self.value)
|
||||
gen_marker_index = progress.find(
|
||||
f'"{self.generation_marker}"'
|
||||
)
|
||||
if gen_marker_index != -1:
|
||||
progress = progress[:gen_marker_index]
|
||||
else:
|
||||
raise ValueError("Failed to find generation marker")
|
||||
|
||||
prompt = template.format(
|
||||
prompt=self.prompt,
|
||||
schema=json.dumps(self.json_schema),
|
||||
progress=progress,
|
||||
)
|
||||
|
||||
return prompt
|
||||
|
||||
def __call__(self) -> Dict[str, Any]:
|
||||
self.value = {}
|
||||
generated_data = self.generate_object(
|
||||
self.json_schema["properties"], self.value
|
||||
)
|
||||
return generated_data
|
||||
@ -1,109 +0,0 @@
|
||||
from swarms.utils.auto_download_check_packages import (
|
||||
auto_check_and_download_package,
|
||||
)
|
||||
from typing import Any
|
||||
|
||||
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
auto_check_and_download_package(
|
||||
"torch", package_manager="pip", upgrade=True
|
||||
)
|
||||
import torch
|
||||
|
||||
try:
|
||||
import transformers
|
||||
except ImportError:
|
||||
auto_check_and_download_package(
|
||||
"transformers", package_manager="pip", upgrade=True
|
||||
)
|
||||
import transformers
|
||||
|
||||
|
||||
class StringStoppingCriteria(transformers.StoppingCriteria):
|
||||
def __init__(
|
||||
self, tokenizer: Any, prompt_length: int # type: ignore
|
||||
):
|
||||
self.tokenizer = tokenizer
|
||||
self.prompt_length = prompt_length
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
input_ids: torch.LongTensor, # type: ignore
|
||||
_,
|
||||
) -> bool:
|
||||
if len(input_ids[0]) <= self.prompt_length:
|
||||
return False
|
||||
|
||||
last_token_id = input_ids[0][-1]
|
||||
last_token = self.tokenizer.decode(
|
||||
last_token_id, skip_special_tokens=True
|
||||
)
|
||||
|
||||
result = '"' in last_token
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class NumberStoppingCriteria(transformers.StoppingCriteria):
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: Any, # type: ignore
|
||||
prompt_length: int,
|
||||
precision: int = 3,
|
||||
):
|
||||
self.tokenizer = tokenizer
|
||||
self.precision = precision
|
||||
self.prompt_length = prompt_length
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
input_ids: torch.LongTensor, # type: ignore
|
||||
scores: torch.FloatTensor, # type: ignore
|
||||
) -> bool:
|
||||
decoded = self.tokenizer.decode(
|
||||
input_ids[0][self.prompt_length :],
|
||||
skip_special_tokens=True,
|
||||
)
|
||||
|
||||
if decoded.count(".") > 1:
|
||||
return True
|
||||
|
||||
if (
|
||||
decoded.count(".") == 1
|
||||
and len(decoded.strip().split(".")[1]) > self.precision
|
||||
):
|
||||
return True
|
||||
|
||||
if (
|
||||
len(decoded) > 1
|
||||
and any(c.isdigit() for c in decoded)
|
||||
and decoded[-1] in [" ", "\n"]
|
||||
):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
class OutputNumbersTokens(transformers.LogitsWarper):
|
||||
def __init__(self, tokenizer: transformers.PreTrainedTokenizer, prompt: str): # type: ignore
|
||||
self.tokenizer = tokenizer
|
||||
self.tokenized_prompt = tokenizer(prompt, return_tensors="pt")
|
||||
vocab_size = len(tokenizer)
|
||||
self.allowed_mask = torch.zeros(vocab_size, dtype=torch.bool)
|
||||
|
||||
for _, token_id in tokenizer.get_vocab().items():
|
||||
token_str = tokenizer.decode(token_id).strip()
|
||||
|
||||
if token_str == "" or (
|
||||
all(c.isdigit() or c == "." for c in token_str)
|
||||
and token_str.count(".") <= 1
|
||||
):
|
||||
self.allowed_mask[token_id] = True
|
||||
|
||||
def __call__(self, _, scores):
|
||||
mask = self.allowed_mask.expand_as(scores)
|
||||
scores[~mask] = -float("inf")
|
||||
|
||||
return scores
|
||||
Loading…
Reference in new issue