simplifies to_word

simplifies the definition of `to_word` with `np.random.choice`
lucasjinreal · Jun 4, 2018 · e24303d · e24303d
1 parent fc36b63
commit e24303d
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/compose_poem.py b/compose_poem.py
@@ -30,12 +30,13 @@
 
 
 def to_word(predict, vocabs):
-    t = np.cumsum(predict)
-    s = np.sum(predict)
-    sample = int(np.searchsorted(t, np.random.rand(1) * s))
+    predict = predict[0]       
+    predict /= np.sum(predict)
+    sample = np.random.choice(np.arange(len(predict)), p=predict)
     if sample > len(vocabs):
-        sample = len(vocabs) - 1
-    return vocabs[sample]
+        return vocabs[-1]
+    else:
+        return vocabs[sample]
 
 
 def gen_poem(begin_word):