aitool.py

from transformers import Tool


class AITranslationTool(Tool):
    name = "model_translation"
    description = (
        "This is a tool that translates "
    )

    inputs = ["text"]
    outputs = ["text"]

    def __call__(self, task: str):
        model = "facebook/nllb-200-distilled-600M"
        return model.id