Tokenizers: How to get representation for a single word form subwords

Hi!

I am using a Roberta model to encode words and then finetune on a downstream task for which I need to find relations between words, not subwords. I would like to represent words as either the mean subword embedding for subwords comprising that word, or as just the first subword.

How can I efficiently get a handle on which subwords comprise a single token, when my input is an entire sentence, from the output of RobertaModel, so I can then compose them in some way? Thanks!