Refactored out sampling from log probs

Easier to see what we're doing. I haven't checked performance.