File size: 385 Bytes
0d55eff
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
from transformers import Tool
import pandas as pd

class PandasDataFrameTool(Tool):
    name = "pandas_dataframe_tool"
    description = "This tool counts the number of rows in a Pandas DataFrame."

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

    def __call__(self, dataframe: pd.DataFrame):
        num_rows = len(dataframe)
        return f"The DataFrame has {num_rows} rows."