Update retriever.py

retriever.py

@@ -2,7 +2,7 @@ from smolagents import Tool
 
 class FrugalAI_methods(Tool):
     name = "FrugalAI_methods"
-    description = "Retrieves methods for model frugalization."
+    description = "Retrieves methods for model frugalization. It will return ideas to frugalize a code, please use it."
     inputs = {
         "method": {
             "type": "string",
@@ -10,25 +10,31 @@ class FrugalAI_methods(Tool):
         }
     }
     output_type = "string"
-
-    def pruning(self,method: str):
-        """
-        Optimizes models by removing unnecessary components, such as certain weights in a neural network. 
-        This function demonstrates how to apply pruning.
-        """
-        model = apply_pruning(model, amount=0.3)
-        code = "model = apply_pruning(model, amount=0.3)"  
-        return (
+    
+    def forward(self, method):
+        ideas=[]
+        ideas.append(pruning())
+        ideas.append(quantization())
+        return ideas
+        
+def pruning(self,method: str):
+    """
+    Optimizes models by removing unnecessary components, such as certain weights in a neural network. 
+    This function demonstrates how to apply pruning.
+    """
+    model = apply_pruning(model, amount=0.3)
+    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)."
         )
 
-    def quantization(self, method: str):
-        """
-        Converts high-precision weights into lower-precision one to reduce cost. 
-        """
-        code = "model = torch.quantization.quantize_dynamic(model, dtype=torch.qint8)"
-        return (
+def quantization(self, method: str):
+    """
+    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}."
         )