107 lines
2.7 KiB
Python
107 lines
2.7 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Example of model selection using agentscope tuner."""
|
|
|
|
import os
|
|
import logging
|
|
from typing import Dict, Any
|
|
from agentscope.agent import ReActAgent
|
|
from agentscope.formatter import OpenAIChatFormatter
|
|
from agentscope.message import Msg
|
|
from agentscope.model import DashScopeChatModel
|
|
from agentscope.tuner import DatasetConfig
|
|
from agentscope.tuner import WorkflowOutput
|
|
from agentscope.tuner.model_selection import select_model
|
|
from agentscope.tuner.model_selection import avg_token_consumption_judge
|
|
|
|
|
|
# Configure logging to show INFO level messages
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
|
|
# Initialize models for selection
|
|
models = [
|
|
DashScopeChatModel(
|
|
"qwen-turbo",
|
|
api_key=os.environ.get("DASHSCOPE_API_KEY", ""),
|
|
max_tokens=512,
|
|
),
|
|
DashScopeChatModel(
|
|
"qwen-plus",
|
|
api_key=os.environ.get("DASHSCOPE_API_KEY", ""),
|
|
max_tokens=512,
|
|
),
|
|
DashScopeChatModel(
|
|
"qwen-max",
|
|
api_key=os.environ.get("DASHSCOPE_API_KEY", ""),
|
|
max_tokens=512,
|
|
),
|
|
]
|
|
|
|
|
|
async def workflow(
|
|
task: Dict[str, Any],
|
|
model: Any,
|
|
) -> WorkflowOutput:
|
|
"""A workflow function using the ReAct agent to solve tasks.
|
|
|
|
Args:
|
|
task (Dict[str, Any]): The task to be solved.
|
|
model: The model to use for the agent.
|
|
|
|
Returns:
|
|
WorkflowOutput: The workflow output containing the agent's response.
|
|
"""
|
|
agent = ReActAgent(
|
|
name="math_solver",
|
|
sys_prompt="You are a helpful math problem solving agent.",
|
|
model=model,
|
|
formatter=OpenAIChatFormatter(),
|
|
)
|
|
|
|
# Extract question from task
|
|
question = (
|
|
task.get("question", "") if isinstance(task, dict) else str(task)
|
|
)
|
|
|
|
# Create a message with the question
|
|
msg = Msg(name="user", content=question, role="user")
|
|
|
|
# Get response from the agent
|
|
response = await agent.reply(msg=msg)
|
|
|
|
return WorkflowOutput(
|
|
response=response,
|
|
)
|
|
|
|
|
|
async def main() -> None:
|
|
"""Main entry point to run model selection example.
|
|
|
|
This function selects the best model based on
|
|
token consumption, and prints the results.
|
|
"""
|
|
# Configure the GSM8K dataset
|
|
dataset_config = DatasetConfig(
|
|
path="openai/gsm8k",
|
|
name="main",
|
|
split="test",
|
|
total_steps=20, # Limit for testing purposes
|
|
)
|
|
|
|
# Perform model selection
|
|
best_model, metrics = await select_model(
|
|
workflow_func=workflow,
|
|
judge_func=avg_token_consumption_judge,
|
|
train_dataset=dataset_config,
|
|
candidate_models=models,
|
|
)
|
|
|
|
print(f"Selected best model: {best_model.model_name}")
|
|
print(f"Metrics: {metrics}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import asyncio
|
|
|
|
asyncio.run(main())
|