fix generating with all tags if all probability < 0

toshiaki1729 · Dec 29, 2022 · 238318c · 238318c
1 parent 362da88
commit 238318c
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/scripts/t2p/prompt_generator/wd_like.py b/scripts/t2p/prompt_generator/wd_like.py
@@ -130,7 +130,7 @@ def __call__(self, text: str, text_neg: str, neg_weight: float, opts: pgen.Gener
                 probs_np = probs_cpu.detach().numpy()
 
                 if np.count_nonzero(probs_np) <= opts.n:
-                    results = tags_np
+                    results = np.random.choice(a=tags_np, size=opts.n, replace=False)
                 else:
                     results = np.random.choice(a=tags_np, size=opts.n, replace=False, p=probs_np)
             else:
@@ -150,7 +150,7 @@ def __call__(self, text: str, text_neg: str, neg_weight: float, opts: pgen.Gener
                 probs_np = np.nan_to_num(probs_np)
 
                 if np.count_nonzero(probs_np) <= opts.n:
-                    results = tags_np
+                    results = np.random.choice(tags_np, opts.n, replace=False)
                 else:
                     results = np.random.choice(tags_np, opts.n, replace=False, p=probs_np)
             else:
@@ -176,7 +176,7 @@ def __call__(self, text: str, text_neg: str, neg_weight: float, opts: pgen.Gener
                 probs_np = np.nan_to_num(probs_np)
 
                 if np.count_nonzero(probs_np) <= opts.n:
-                    results = tags_np
+                    results = np.random.choice(tags_np, opts.n, replace=False)
                 else:
                     results = np.random.choice(tags_np, opts.n, replace=False, p=probs_np)
             else: