agentscope/examples/tuner/model_selection/example_bleu.py

# -*- coding: utf-8 -*-
"""Example of model selection for translation tasks 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 JudgeOutput
from agentscope.tuner import WorkflowOutput
from agentscope.tuner import DatasetConfig
from agentscope.tuner.model_selection import select_model


# Configure logging to show INFO level messages
logging.basicConfig(level=logging.INFO)


# Initialize models for selection
models = [
    DashScopeChatModel(
        "qwen3-max-2025-09-23",
        api_key=os.environ.get("DASHSCOPE_API_KEY", ""),
        max_tokens=1024,
    ),
    DashScopeChatModel(
        "Moonshot-Kimi-K2-Instruct",
        api_key=os.environ.get("DASHSCOPE_API_KEY", ""),
        max_tokens=1024,
    ),
    DashScopeChatModel(
        "MiniMax-M2.1",
        api_key=os.environ.get("DASHSCOPE_API_KEY", ""),
        max_tokens=1024,
    ),
    DashScopeChatModel(
        "deepseek-r1",
        api_key=os.environ.get("DASHSCOPE_API_KEY", ""),
        max_tokens=1024,
    ),
    DashScopeChatModel(
        "glm-4.7",
        api_key=os.environ.get("DASHSCOPE_API_KEY", ""),
        max_tokens=1024,
    ),
]


async def translation_workflow(
    task: Dict[str, Any],
    model: Any,
) -> WorkflowOutput:
    """A workflow function using the ReAct agent to perform translation tasks.

    Args:
        task (Dict[str, Any]): The translation task
        containing source text and target language.
        model: The model to use for the agent.

    Returns:
        WorkflowOutput: The workflow output containing the translated text.
    """
    agent = ReActAgent(
        name="translator",
        sys_prompt=(
            "You are a helpful translation agent."
            "Only output the translated text."
        ),
        model=model,
        formatter=OpenAIChatFormatter(),
    )

    # Extract source text and target language from task
    source_text = (
        task.get("question", "") if isinstance(task, dict) else str(task)
    )

    # Create a message with the translation request
    prompt = (
        f"Translate following text between English and Chinese: {source_text}"
    )
    msg = Msg(name="user", content=prompt, role="user")

    # Get response from the agent
    response = await agent.reply(msg=msg)

    return WorkflowOutput(
        response=response,
    )


async def bleu_judge(
    task: Dict[str, Any],
    response: Any,
) -> JudgeOutput:
    """A judge function to calculate BLEU score for translation quality.

    Args:
        task (Dict[str, Any]): The task information.
        response (Any): A composite dict containing the workflow response
            and metrics.

    Returns:
        JudgeOutput: The BLEU score and other metrics.
    """
    # Lazy import to follow the requirement
    import sacrebleu

    response_str = ""
    if isinstance(response, dict) and "response" in response:
        response_content = response["response"]
        if hasattr(response_content, "content"):
            # Handle the response structure
            if isinstance(response_content.content, list):
                for content_item in response_content.content:
                    if (
                        isinstance(content_item, dict)
                        and "text" in content_item
                    ):
                        response_str += content_item["text"]
                    elif hasattr(content_item, "text"):
                        response_str += content_item.text
            else:
                response_str = str(response_content.content)
    else:
        raise ValueError("Response must be a dict with 'response' key")

    # Extract reference translation
    reference_translation = (
        task.get("answer", "") if isinstance(task, dict) else ""
    )

    # Load the BLEU metric
    ref = reference_translation.strip()
    pred = response_str.strip()

    bleu_score = sacrebleu.sentence_bleu(pred, [ref])

    # Return the judge output with the BLEU score as reward
    return JudgeOutput(
        reward=bleu_score.score,
        metrics={
            "bleu": bleu_score.score / 100,
            "brevity_penalty": bleu_score.bp,
            "ratio": bleu_score.ratio,
        },
    )


async def main() -> None:
    """Main entry point to run model selection example.

    This function selects the best model based on
    bleu score, and prints the results.
    """
    # Define the translation benchmark dataset using DatasetConfig
    dataset_config = DatasetConfig(
        path=os.path.join(
            os.path.dirname(__file__),
            "translate_data",
        ),  # Path to the local JSON dataset
        split="test",
    )

    # Perform model selection using the local translation benchmark dataset
    best_model, metrics = await select_model(
        workflow_func=translation_workflow,
        judge_func=bleu_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())