Make sampling with bias numerically stable (#90)

* Update sampling.py

Remove a slow for loop on logit bias. Make the numpy re-softmax operation numerically stable.

* Update sampling.py

rwkv/sampling.py

@@ -18,9 +18,11 @@ def sample_probs(probs: np.ndarray, temperature: float = 1.0, top_p: float = 0.8
     if logit_bias is not None:
         logits = np.log(probs)
 
-        for token in logit_bias.keys():
-            logits[token] += logit_bias[token]
-
+        if len(logit_bias) > 0:
+            ids, values = zip(*logit_bias.items())
+            logits[list(ids)] += values
+
+        logits -= logits.max(axis=-1, keepdims=True)
         probs = np.exp(logits) / np.sum(np.exp(logits))
 
     if temperature == 0.0: