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