Apply zero-shot transformer model to each row and create new column(s) in pandas for the appropriate label (custom function and .apply)

Viewed 28

Right now I have this Huggingface transformer pipeline that does zero-shot classification. I want to apply it to a open-answer column from a survey dataset, where I run the model rowwise and create a new column called "theme" (the label with the highest probability score) and a second column called "prob" that lists the probability score for the label that got added to the "theme" column.

As of now, I have this mock up dataset:

                                             open_text  col_2  col_3
0  The way he threw that 3-pointer was incredible       NaN    NaN
1  On election day, people tend to queue for way ...    NaN    NaN
2  She did not order it because she was already full    NaN    NaN
3  He enjoyed his hot-dog watching the Lakers game      NaN    NaN

I have written this custom function:

def zeroshotPipeline(text):
    input_ids = text
    candidate_label = ['basketball', 'politics', 'food']
    template = "This example is {}"
    results = classifier(input_ids, 
                         candidate_label,
                        # multi_label = True,
                         hypothesis_template = template)
    score_id = np.argmax(results["scores"])
    final_label = results["labels"][score_id]
    prob = results["scores"][score_id]
    return final_label, prob


df["theme"] = ""
df["prob"] = np.nan

df['theme'] = df["open_text"].apply(zeroshotPipeline)
print(df)

                                           open_text  ...                            theme
0     The way he threw that 3-pointer was incredible  ...  (politics, 0.12852472066879272)
1  On election day, people tend to queue for way ...  ...   (politics, 0.9359661340713501)
2  She did not order it because she was already full  ...       (food, 0.9898027181625366)
3    He enjoyed his hot-dog watching the Lakers game  ...         (food, 0.99793541431427)

As you can see, when I print df the column "theme" consists of the label and the probability, separated by a comma. I would like these to be in each of their column. How is this done?

Additionally, is there a way to do it where I add multi_label = True inside the function, and add the columns theme_1, theme_2 (..) according to those labels that have a probability of say > 0.60? So in the example above, the results for the 4th row (when setting multi-label = True):

{'sequence': 'He enjoyed his hot-dog watching the Lakers game', 'labels': ['food', 'basketball', 'politics'], 'scores': [0.99793541431427, 0.9612331390380859, 0.01709340512752533]}

Then I would like to include two new columns (theme_1 and theme_2) where the labels "food" and "basketball" are included, followed by two columns (prob_1 and prob_2) where their corresponding probabilities are the values.

I'm relatively new to python (coming from R), so this rather "simple" wrangling is just something I have trouble achieving in python.

Thanks

1 Answers

first split the column that has the tuple into a new dataframe with just that column. Do something like this:

split_df = pd.DataFrame(df['theme'].tolist(), columns=['theme', 'score'])

then delete the column you want to change from the initial dataframe like this:

df.drop('theme')

Then combine the two data frames together like this:

df = pd.concat([df, split_df], axis=1)

and that should do it for you.

Related