zeno_ml.py

import pickle


from datasets import load_dataset
from zeno import (
    DistillReturn,
    MetricReturn,
    ModelReturn,
    ZenoOptions,
    distill,
    metric,
    model,
)

@model
def get_model(model_name):


    def pred(df, ops: ZenoOptions):
        res = []
        #print(df)
        #print(ops)
        for data_columns in df[ops.data_column]:
            ## return where the data column exist in the dataframe column Question by the index
            df_index = df[df["Question"] == data_columns].index[0]
            res.append(df["Answering"][df_index])
        return ModelReturn(model_output=res)

    return pred