Efficiently use a fill-mask pipeline on a large pretrained transformers model to categorise a list of words

categorization_fill_mask.py

from transformers import pipeline

pipe = pipeline("fill-mask", model="roberta-large", top_k=1)

def get_topic_using_fill_mask(words):
    results = pipe(f"{words}. \n These words are about <mask>.")

    return results[0]["token_str"].strip()
    
get_topic_using_fill_mask(["terra, mars, moon, saturn, jupyter"])

# prints "astronomy"