Hi, I'm learning how to use GPT-J-6B. I am wondering if there is a way to do something like autocompletion with it? Based on input, generate a set of expected next words with probabilities?

I see many tutorials, but they all use it to just generate lots of text. But I'm looking for it to generate a tree, not just the DFS path. If someone can help point me to a specific function in the codebase, I would appreciate it.