Upload generate.py with huggingface_hub

generate.py

@@ -37,6 +37,9 @@ def main():
     parser.add_argument("--repetition-penalty", type=float, default=1.0)
     parser.add_argument("--seed", type=int, default=None)
     parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"])
+    parser.add_argument("--guard", default=None, help="file kamus (satu kata per baris); pilih kandidat dengan ejaan terbaik")
+    parser.add_argument("--guard-tries", type=int, default=5, help="jumlah kandidat saat --guard aktif")
+    parser.add_argument("--guard-min", type=float, default=0.6, help="rasio kata dikenal minimum (info saja)")
     args = parser.parse_args()
 
     if args.seed is not None:
@@ -46,17 +49,45 @@ def main():
     model, tokenizer = load_model(args.ckpt)
     model = model.to(device)
 
+    wordset = None
+    if args.guard:
+        from indigo.common import load_wordlist, word_known_ratio
+
+        wordset = load_wordlist(args.guard)
+        print(f"[guard] kamus: {len(wordset):,} kata | target rasio >= {args.guard_min:.0%}")
+
     ids = tokenizer.encode(args.prompt) or [0]
     idx = torch.tensor([ids], dtype=torch.long, device=device)
-    out = model.generate(
-        idx,
-        args.max_new,
-        temperature=args.temperature,
-        top_k=args.top_k,
-        top_p=args.top_p,
-        repetition_penalty=args.repetition_penalty,
-    )
-    print(tokenizer.decode(out[0].tolist()))
+
+    def sample():
+        out = model.generate(
+            idx,
+            args.max_new,
+            temperature=args.temperature,
+            top_k=args.top_k,
+            top_p=args.top_p,
+            repetition_penalty=args.repetition_penalty,
+        )
+        text = tokenizer.decode(out[0].tolist())
+        ratio = word_known_ratio(text, wordset) if wordset else 1.0
+        return text, ratio
+
+    if wordset is None:
+        text, _ = sample()
+        print(text)
+        return
+
+    best_text, best_ratio = "", -1.0
+    for t in range(args.guard_tries):
+        torch.manual_seed((args.seed or 0) + t * 1013)
+        text, ratio = sample()
+        mark = f"  [kandidat {t + 1}: {ratio:.0%}]"
+        if ratio > best_ratio:
+            best_text, best_ratio = text, ratio
+        if best_ratio >= args.guard_min:
+            break
+    print(best_text)
+    print(f"[guard] rasio kata dikenal: {best_ratio:.0%}")
 
 
 if __name__ == "__main__":