from smolagents import Tool

class FrugalAI_methods(Tool):
    name = "FrugalAI_methods"
    description = "Retrieves methods for model frugalization. It will return ideas to frugalize a code, please use it."
    inputs = {
        "method": {
            "type": "string",
            "description": "Always set it to 'Frugal' or it will raise an error."
        }
    }
    output_type = "string"
    
    def forward(self, method):
        ideas=[]
        ideas.append(pruning())
        ideas.append(quantization())
        return ideas
        
def pruning():
    """
    Optimizes models by removing unnecessary components, such as certain weights in a neural network. 
    This function demonstrates how to apply pruning.
    """
    code = "model = apply_pruning(model, amount=0.3)"
    return (
            f"To apply pruning to a model, use the following code snippet: {code}. "
            f"You should adapt it to your actual implementation. In particular, the 'amount' parameter "
            f"can be increased or decreased depending on the initial number of weights and the complexity of your use case (minimu value: 0, maximum value: 1). You should give these information to the manager, you can not prune the model yourself."
        )

def quantization():
    """
    Converts high-precision weights into lower-precision one to reduce cost. 
    """
    code = "model = torch.quantization.quantize_dynamic(model, dtype=torch.qint8)"
    return (
            f"To apply quantization to a model, use the following code snippet: {code}. You should give these information to the manager, you can not quantize the model yourself."
        )